from __future__ import annotations
from pathlib import Path
import json
import joblib
import numpy as np
import xarray as xr
from .io_utils import load_grid, write_dataset_safe, TARGET_TO_OBS
from .time_utils import split_masks
from .transform import monthly_climatology, target_transform, inverse_target, obs_physical, raw_physical
from .features import build_feature_cube, select_samples, select_samples_with_aux, predict_to_grid, spatial_masks
from .models import fit_with_optional_search
from .tasks import task_label


def _indices(mask: np.ndarray, months: np.ndarray, month: int) -> np.ndarray:
    return np.where(mask & (months == int(month)))[0]


def _fit_predict_month(cfg, gd, clim, task, fit_mask, pred_idx, month: int, final_fit: bool):
    target = task['target']
    feature_vars = cfg['members']['feature_sets'][target][task['feature_set']]
    include_lat_lon = bool(cfg['training'].get('include_lat_lon', True))
    include_reg = bool(cfg['training'].get('include_region_feature', False))
    precip_floor = float(cfg['training'].get('precip_clim_floor', 1.0))
    fit_idx = _indices(fit_mask, gd.months, month)
    if len(pred_idx) == 0:
        return None, {}
    X_fit_cube, feature_names = build_feature_cube(gd.ds, clim, feature_vars, fit_idx, gd.months, gd.lat2d, gd.lon2d, gd.reg, include_lat_lon, include_reg, precip_floor)
    y_fit_cube = target_transform(gd.ds, clim, target, fit_idx, gd.months, precip_floor)
    physical_scale_cube = None
    if target == 'pr':
        # For hyperparameter tuning, score pr in physical units after inverse transform:
        # pr_bc = max(pr_factor_hat * pr_o_cli, 0).  The appended scale is used only
        # by the CV scorer; it is not used as a predictor.
        scale2d = clim['pr_o'][month-1].astype(np.float32)
        physical_scale_cube = np.broadcast_to(scale2d[None, :, :], y_fit_cube.shape).astype(np.float32)
    X_pred_cube, _ = build_feature_cube(gd.ds, clim, feature_vars, pred_idx, gd.months, gd.lat2d, gd.lon2d, gd.reg, include_lat_lon, include_reg, precip_floor)
    pred_tr = np.full(X_pred_cube.shape[:3], np.nan, dtype=np.float32)
    model_info = {'feature_names': feature_names, 'month': month, 'regions': {}}
    for region_label, smask in spatial_masks(gd.valid_mask, gd.reg, task['spatial_mode'], cfg['grid'].get('valid_regions', [1,2,3,4,5,6,7])):
        X, y, w_phys = select_samples_with_aux(X_fit_cube, y_fit_cube, smask, physical_scale_cube)
        if len(y) < int(cfg['training'].get('min_train_samples', 30)):
            model_info['regions'][region_label] = {'status': 'skip_small_sample', 'n': int(len(y))}
            continue
        model, info = fit_with_optional_search(
            X, y, task['ai_model'], cfg,
            physical_scale=w_phys if target == 'pr' else None,
            clip_nonnegative=bool(cfg['training'].get('clip_precipitation_to_zero', True)),
        )
        p = predict_to_grid(model, X_pred_cube, smask)
        pred_tr[:, smask] = p[:, smask]
        model_info['regions'][region_label] = {'status': 'fit', 'n': int(len(y)), **info}
        model_info['regions'][region_label]['model'] = model if final_fit else None
    pred = inverse_target(pred_tr, clim, target, pred_idx, gd.months, clip_pr=bool(cfg['training'].get('clip_precipitation_to_zero', True)))
    pred[:, ~gd.valid_mask] = np.nan
    return pred, model_info


def run_bc_task(cfg: dict, task: dict, dirs: dict[str, Path]) -> dict:
    label = task_label(task)
    done = dirs['done'] / f'bc_{label}.done'
    eval_out = dirs['evaluation_members'] / task['target'] / f'{label}.nc'
    prod_out = dirs['production_members'] / task['target'] / f'{label}.nc'
    final_art = dirs['final_models'] / task['target'] / f'{label}.joblib'
    eval_art = dirs['models'] / 'evaluation' / task['target'] / f'{label}.json'
    for p in [eval_out, prod_out, final_art, eval_art]:
        p.parent.mkdir(parents=True, exist_ok=True)
    if done.exists() and eval_out.exists() and prod_out.exists() and final_art.exists():
        return {'status': 'skipped', 'label': label}
    gd = load_grid(cfg, task['source_model'], int(task['lead']))
    masks = split_masks(cfg, gd.time, int(task['lead']))
    ref_mask = masks['clim']
    clim = monthly_climatology(gd.ds, gd.months, ref_mask, precip_floor=float(cfg['training'].get('precip_clim_floor', 1.0)))
    eval_mask = masks['train'] | masks['valid']
    eval_idx = np.where(eval_mask)[0]
    full_idx = np.arange(gd.ds.sizes['mon'])
    nlat, nlon = gd.ds.sizes['lat'], gd.ds.sizes['lon']
    bc_eval = np.full((len(eval_idx), nlat, nlon), np.nan, dtype=np.float32)
    bc_prod = np.full((len(full_idx), nlat, nlon), np.nan, dtype=np.float32)
    eval_info = {'label': label, 'task': dict(task), 'months': {}}
    final_info = {'label': label, 'task': dict(task), 'months': {}}
    final_models = {'label': label, 'task': dict(task), 'climatology_reference': str(ref_mask.sum()), 'monthly_models': {}}
    for month in range(1,13):
        # Evaluation models: train on training split, predict training+validation months.
        mi_eval = np.where(gd.months[eval_idx] == month)[0]
        if len(mi_eval):
            pred, info = _fit_predict_month(cfg, gd, clim, task, masks['train'], eval_idx[mi_eval], month, final_fit=False)
            if pred is not None:
                bc_eval[mi_eval] = pred
            eval_info['months'][str(month)] = info
        # Final production models: refit on train+validation/full valid split and predict all 480 months.
        pi = np.where(gd.months[full_idx] == month)[0]
        predp, infof = _fit_predict_month(cfg, gd, clim, task, masks['full'], full_idx[pi], month, final_fit=True)
        if predp is not None:
            bc_prod[pi] = predp
        # Save model objects separately by re-fitting once more per region would be expensive; current info does not hold models in JSON.
        final_info['months'][str(month)] = {k:v for k,v in infof.items() if k != 'regions'}
        final_models['monthly_models'][str(month)] = infof
    # strip non-serializable model objects from JSON but keep them in joblib
    def strip_models(obj):
        if isinstance(obj, dict):
            return {k: strip_models(v) for k,v in obj.items() if k != 'model'}
        if isinstance(obj, list):
            return [strip_models(v) for v in obj]
        try:
            json.dumps(obj)
            return obj
        except Exception:
            return str(obj)
    obs_eval = obs_physical(gd.ds, task['target'], eval_idx)
    raw_eval = raw_physical(gd.ds, task['target'], eval_idx)
    split_flag = np.where(masks['train'][eval_idx], 1, np.where(masks['valid'][eval_idx], 2, 0)).astype(np.int16)
    ds_eval = xr.Dataset(coords={'time': gd.time[eval_idx], 'lat': gd.ds['lat'].values, 'lon': gd.ds['lon'].values}, data_vars={
        'obs': (('time','lat','lon'), obs_eval.astype(np.float32)),
        'raw': (('time','lat','lon'), raw_eval.astype(np.float32)),
        'bc': (('time','lat','lon'), bc_eval.astype(np.float32)),
        'split': (('time',), split_flag),
    }, attrs={'label': label, **{k: str(v) for k,v in task.items()}, 'split_1': 'training', 'split_2': 'validation'})
    ds_prod = xr.Dataset(coords={'time': gd.time, 'lat': gd.ds['lat'].values, 'lon': gd.ds['lon'].values}, data_vars={
        'bc': (('time','lat','lon'), bc_prod.astype(np.float32)),
        'obs': (('time','lat','lon'), obs_physical(gd.ds, task['target'], full_idx).astype(np.float32)),
    }, attrs={'label': label, **{k: str(v) for k,v in task.items()}, 'purpose': 'final production refit on full train+validation split for GD2'})
    write_dataset_safe(ds_eval, eval_out)
    write_dataset_safe(ds_prod, prod_out)
    eval_art.write_text(json.dumps(strip_models(eval_info), indent=2), encoding='utf-8')
    joblib.dump(final_models, final_art, compress=3)
    done.write_text(json.dumps({'status':'done','label':label,'evaluation':str(eval_out),'production':str(prod_out),'artifact':str(final_art)}, indent=2), encoding='utf-8')
    return {'status':'done','label':label}