from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple

import numpy as np
import xarray as xr

from .data_utils import ALL_DATA_VARIABLES, TARGET_TO_MODEL, compute_train_climatology, load_grid_data, raw_for_var, year_range_mask, target_training_transform, inverse_target_transform
from .features import build_feature_cube, predict_to_grid, select_samples
from .models import make_model
from .train_eval import _spatial_masks, task_label, _target_units


def run_fullperiod_bc_task(cfg: Dict[str, Any], task: Dict[str, Any], dirs: Dict[str, Path]) -> Dict[str, Path]:
    """Train one BC member on final_fit_years and predict the full 480-month period.

    Output is intended for SPEI/hindcast construction, not just test-period evaluation.
    """
    label = task_label(task)
    out_file = dirs['corrected_full'] / f'bc_full_{label}.nc'
    done_file = dirs['done'] / f'bc_full_{label}.done'
    if done_file.exists() and out_file.exists() and not cfg.get('overwrite', False):
        return {'status': 'skipped', 'label': label, 'full_bc': out_file}

    gd = load_grid_data(
        data_dir=cfg['data_dir'],
        mask_file=cfg['mask_file'],
        file_pattern=cfg.get('file_pattern', 'lead{lead}_{source_model}_01_BL_Test_VN.nc'),
        source_model=task['source_model'],
        lead=int(task['lead']),
        year_start=int(cfg.get('year_start', 1982)),
        missing_value=float(cfg.get('missing_value', -99.0)),
    )
    clim = compute_train_climatology(gd.ds, gd.years, cfg.get('climatology_years', cfg['fit_years']), list(dict.fromkeys(ALL_DATA_VARIABLES)))
    fit_mask_all = year_range_mask(gd.years, cfg['fit_years'])
    ntime, nlat, nlon = gd.ds.sizes['mon'], gd.ds.sizes['lat'], gd.ds.sizes['lon']
    pred = np.full((ntime, nlat, nlon), np.nan, dtype=np.float32)
    raw = np.full_like(pred, np.nan)
    obs = np.full_like(pred, np.nan)

    target = task['target']
    method = task['method'].lower()
    predictor_set = task['predictor_set']
    spatial_mode = task['spatial_mode']
    direct_model_var = TARGET_TO_MODEL[target]

    for month in range(1, 13):
        fit_idx = np.where(fit_mask_all & (gd.months == month))[0]
        all_idx = np.where(gd.months == month)[0]
        if len(all_idx) == 0:
            continue
        obs_m = raw_for_var(gd.ds, target, all_idx).astype(np.float32)
        raw_m = raw_for_var(gd.ds, direct_model_var, all_idx).astype(np.float32)
        obs[all_idx] = obs_m
        raw[all_idx] = raw_m
        if method == 'raw':
            pred_m = raw_m.copy()
            if target == 'pr_o' and cfg.get('clip_precipitation_to_zero', True):
                pred_m = np.where(np.isfinite(pred_m), np.maximum(pred_m, 0.0), np.nan).astype(np.float32)
            pred[all_idx] = pred_m
            continue
        target_fit_anom = target_training_transform(gd.ds, clim, target, fit_idx, month)
        feat_fit, _ = build_feature_cube(
            gd.ds, clim, target, predictor_set, fit_idx, month, gd.lat2d, gd.lon2d, gd.reg, spatial_mode,
            include_lat_lon=cfg.get('include_lat_lon', True), include_region_feature=cfg.get('include_region_feature', True),
        )
        feat_all, _ = build_feature_cube(
            gd.ds, clim, target, predictor_set, all_idx, month, gd.lat2d, gd.lon2d, gd.reg, spatial_mode,
            include_lat_lon=cfg.get('include_lat_lon', True), include_region_feature=cfg.get('include_region_feature', True),
        )
        pred_month = np.full((len(all_idx), nlat, nlon), np.nan, dtype=np.float32)
        for region_label, smask in _spatial_masks(gd.valid_mask, gd.reg, spatial_mode, cfg.get('valid_regions', [1,2,3,4,5,6,7])):
            X_fit, y_fit, *_ = select_samples(feat_fit, target_fit_anom, smask)
            if len(y_fit) < int(cfg.get('min_train_samples', 20)):
                continue
            model = make_model(method, cfg, random_seed=int(cfg.get('random_seed', 42)))
            model.fit(X_fit, y_fit)
            pred_anom = predict_to_grid(model, feat_all, smask)
            pred_raw = inverse_target_transform(pred_anom, clim, target, month, clip_precipitation_to_zero=cfg.get('clip_precipitation_to_zero', True))
            pred_month[:, smask] = pred_raw[:, smask]
        if target == 'pr_o' and cfg.get('clip_precipitation_to_zero', True):
            pred_month = np.where(np.isfinite(pred_month), np.maximum(pred_month, 0.0), np.nan).astype(np.float32)
        pred[all_idx] = pred_month

    pred[:, ~gd.valid_mask] = np.nan
    raw[:, ~gd.valid_mask] = np.nan
    obs[:, ~gd.valid_mask] = np.nan
    ds = xr.Dataset(
        coords={'time': gd.time, 'lat': gd.ds['lat'].values, 'lon': gd.ds['lon'].values},
        data_vars={
            'obs': (('time','lat','lon'), obs),
            'raw': (('time','lat','lon'), raw),
            'bc': (('time','lat','lon'), pred),
        },
        attrs={**{k: str(v) for k, v in task.items()}, 'fit_years': f"{cfg['fit_years'][0]}-{cfg['fit_years'][1]}", 'units': _target_units(target)},
    )
    ds['bc'].attrs['long_name'] = 'Full-period bias-corrected prediction'
    ds.to_netcdf(out_file)
    done_file.write_text(json.dumps({'label': label, 'status': 'done', 'file': str(out_file)}, indent=2), encoding='utf-8')
    return {'status': 'done', 'label': label, 'full_bc': out_file}