"""Streaming online fine-tuning: incremental data ingestion, RBF smoothing, and DAP-by-DAP adaptation."""

from __future__ import annotations

import os
import re
from datetime import datetime, timedelta

import numpy as np
import pandas as pd
from sklearn.svm import SVR

from ndvi_forecast.finetune import online_finetune_driver


# ---------------------------------------------------------------------------
# Column normalization
# ---------------------------------------------------------------------------

def normalize_raw_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Normalize raw NDVI column names: 'id' -> 'FID', 6-digit dates -> 8-digit YYYYMMDD."""
    df = df.copy()
    if "id" in df.columns and "FID" not in df.columns:
        df.rename(columns={"id": "FID"}, inplace=True)

    pat = re.compile(r"(\d{6})$")
    renames = {}
    for c in df.columns:
        m = pat.search(str(c))
        if m:
            yy, mm, dd = m.group(1)[:2], m.group(1)[2:4], m.group(1)[4:]
            renames[c] = f"20{yy}{mm}{dd}"
    if renames:
        df.rename(columns=renames, inplace=True)

    df.columns = [str(c) for c in df.columns]
    return df


# ---------------------------------------------------------------------------
# Data sourcing
# ---------------------------------------------------------------------------

def source_data_to_dap(
    df: pd.DataFrame,
    planting_date: str,
    current_dap: int,
    future_window: int,
    date_fmt: str = "%Y%m%d",
) -> tuple[str, pd.DataFrame]:
    """Extract raw NDVI observations up to (current_dap + future_window).

    Returns (latest_date_str, filtered_dataframe).
    """
    ref_dt = datetime.strptime(planting_date, date_fmt) + timedelta(days=int(current_dap))
    cutoff_dt = ref_dt + timedelta(days=int(future_window))

    col_map = {}
    for c in df.columns:
        m = re.search(r"(\d{8})$", str(c))
        if m:
            try:
                col_map[datetime.strptime(m.group(1), date_fmt)] = c
            except Exception:
                pass

    dates = sorted(col_map.keys())
    selected = [d for d in dates if d <= cutoff_dt]
    next_after = next((d for d in dates if d > cutoff_dt), None)
    if next_after:
        selected.append(next_after)

    selected_cols = [col_map[d] for d in selected]
    if selected_cols:
        cols = (["FID"] + selected_cols) if "FID" in df.columns else selected_cols
        result = df[cols].copy()
        latest = selected[-1].strftime(date_fmt)
    else:
        result = df[["FID"]].copy() if "FID" in df.columns else df.iloc[:, :0]
        latest = ref_dt.strftime(date_fmt)

    return latest, result


# ---------------------------------------------------------------------------
# RBF smoothing
# ---------------------------------------------------------------------------

def rbf_smooth(
    ndvi_df: pd.DataFrame,
    planting_date: str,
    rbf_params: dict | None = None,
    date_fmt: str = "%Y%m%d",
) -> tuple[np.ndarray, np.ndarray]:
    """Smooth sparse NDVI observations to a daily DAP grid using RBF-kernel SVR.

    Returns (fid_array, smoothed_ndvi) where smoothed_ndvi has shape (N, max_dap+1).
    """
    if rbf_params is None:
        rbf_params = {"C": 500, "gamma": 0.0001, "epsilon": 1e-6}

    df0 = ndvi_df.copy()
    fid_array = df0["FID"].to_numpy() if "FID" in df0.columns else np.arange(len(df0))

    col_to_date = {}
    for c in df0.columns:
        if c == "FID":
            continue
        m = re.search(r"(\d{8})", str(c))
        if m:
            col_to_date[c] = m.group(1)

    if not col_to_date:
        return fid_array, np.zeros((len(df0), 0), dtype=np.float32)

    cols = sorted(col_to_date.keys(), key=lambda c: col_to_date[c])
    dates = [col_to_date[c] for c in cols]

    plant = datetime.strptime(planting_date, date_fmt)
    daps = np.array([(datetime.strptime(d, date_fmt) - plant).days for d in dates], dtype=int)

    keep = daps >= 0
    cols = [c for c, k in zip(cols, keep) if k]
    daps = daps[keep]

    if len(cols) < 2:
        arr = df0[cols].apply(pd.to_numeric, errors="coerce").to_numpy(np.float32) if cols else np.zeros((len(df0), 0), np.float32)
        return fid_array, arr

    X = daps.reshape(-1, 1)
    t_grid = np.arange(0, int(daps.max()) + 1).reshape(-1, 1)

    Y = df0[cols].apply(pd.to_numeric, errors="coerce").to_numpy(dtype=np.float32)
    smoothed = np.zeros((Y.shape[0], t_grid.shape[0]), dtype=np.float32)

    for i in range(Y.shape[0]):
        row = Y[i]
        mask = np.isfinite(row)
        if mask.sum() < 2:
            smoothed[i, :] = np.nan_to_num(row[mask][-1] if mask.any() else 0.0)
            continue
        svr = SVR(kernel="rbf", **rbf_params)
        svr.fit(X[mask], row[mask])
        smoothed[i, :] = svr.predict(t_grid)

    return fid_array, smoothed


# ---------------------------------------------------------------------------
# Group stats (single-group, FID-aligned)
# ---------------------------------------------------------------------------

def compute_group_stats_at_dap(
    fid_array: np.ndarray,
    ndvi_array: np.ndarray,
    dap_index: int,
) -> np.ndarray:
    """Compute group mean and std at a specific DAP across all samples.

    Returns array of shape (N, 2).
    """
    assert 0 <= dap_index < ndvi_array.shape[1], f"DAP {dap_index} out of range"
    dap_values = ndvi_array[:, dap_index]
    group_mean = float(np.nanmean(dap_values))
    group_std = float(np.nanstd(dap_values, ddof=1))
    N = len(fid_array)
    return np.column_stack([
        np.full(N, group_mean, dtype=np.float32),
        np.full(N, group_std, dtype=np.float32),
    ])


# ---------------------------------------------------------------------------
# Ground truth formatting
# ---------------------------------------------------------------------------

def prepare_ground_truth(df: pd.DataFrame, max_dap: int) -> tuple[np.ndarray, np.ndarray]:
    """Extract FID array and NDVI ground truth matrix from a DataFrame.

    Returns (fid_array, gt_ndvi) where gt_ndvi has shape (N, num_daps).
    """
    dap_cols = sorted([c for c in df.columns if str(c).isdigit() and int(c) <= max_dap], key=int)
    fid = df["FID"].to_numpy() if "FID" in df.columns else np.arange(len(df))
    gt = df[[str(c) for c in dap_cols]].to_numpy(dtype=np.float32)
    return fid, gt


# ---------------------------------------------------------------------------
# Streaming driver
# ---------------------------------------------------------------------------

def run_streaming_finetune(
    model,
    raw_df: pd.DataFrame,
    planting_date: str,
    ground_truth: np.ndarray,
    fid_array: np.ndarray,
    group_stats_dap: int,
    save_dir: str,
    *,
    dap_start: int,
    dap_step: int,
    max_dap: int,
    finetune_window: int,
    device: str = "cpu",
    date_fmt: str = "%Y%m%d",
    rbf_params: dict | None = None,
    ft_kwargs: dict | None = None,
):
    """Incremental online fine-tuning: at each DAP cycle, source new observations,
    smooth to daily grid, and fine-tune the model on the causal window.

    Args:
        model: Pretrained Seq2SeqNDVI model.
        raw_df: Raw satellite NDVI DataFrame with date-named columns.
        planting_date: Planting date string (YYYYMMDD).
        ground_truth: Pre-smoothed NDVI array (N, T) used as labels beyond observed horizon.
        fid_array: Grid cell identifiers aligned with ground_truth rows.
        group_stats_dap: DAP index for computing group statistics.
        save_dir: Directory for output CSVs and checkpoints.
        dap_start: First DAP to fine-tune at.
        dap_step: DAP increment between cycles.
        max_dap: Maximum DAP (sequence length cap).
        finetune_window: Number of future DAPs used as the FT target window.
        device: 'cuda' or 'cpu'.
        date_fmt: Date format string.
        rbf_params: RBF smoothing hyperparameters.
        ft_kwargs: Additional keyword arguments passed to online_finetune_driver.
    """
    ft_kwargs = dict(ft_kwargs or {})
    rbf_params = rbf_params or {"C": 500, "gamma": 0.0001, "epsilon": 1e-6}
    os.makedirs(save_dir, exist_ok=True)

    raw_df = normalize_raw_columns(raw_df)

    GT_T = ground_truth.shape[1]
    end_dap = min(int(max_dap), GT_T)
    print(f"[INIT] N={len(fid_array)} GT_T={GT_T} end_dap={end_dap}")

    def _align_smoothed(fid_smooth, smoothed):
        """Re-align smoothed rows to match fid_array ordering."""
        sm_df = pd.DataFrame(smoothed)
        sm_df.insert(0, "FID", fid_smooth)
        sm_df = sm_df[sm_df["FID"].isin(fid_array)]
        sm_df = sm_df.set_index("FID").loc[fid_array].reset_index()
        return sm_df.drop(columns=["FID"]).to_numpy(dtype=np.float32)

    current_dap = int(dap_start)
    cycle = 0

    while current_dap + int(finetune_window) <= end_dap:
        cycle += 1
        print(f"\n[CYCLE {cycle}] DAP={current_dap}")

        # 1. Source raw observations up to current_dap + finetune_window
        _, raw_slice = source_data_to_dap(
            raw_df, planting_date, current_dap, finetune_window, date_fmt,
        )

        # 2. Smooth to daily DAP grid
        fid_smooth, smoothed = rbf_smooth(raw_slice, planting_date, rbf_params, date_fmt)
        smoothed = _align_smoothed(fid_smooth, smoothed)
        observed_len = smoothed.shape[1]

        if observed_len <= current_dap:
            print(f"[WARN] Smoothed data only to DAP {observed_len - 1}, need {current_dap}. Skipping.")
            current_dap += dap_step
            continue

        # 3. Concatenate: smoothed history + GT labels beyond observed horizon
        T_use = min(end_dap, max_dap)
        ndvi_concat = np.full((len(fid_array), T_use), np.nan, dtype=np.float32)
        hist_T = min(observed_len, T_use)
        ndvi_concat[:, :hist_T] = smoothed[:, :hist_T]
        if hist_T < T_use:
            ndvi_concat[:, hist_T:T_use] = ground_truth[:, hist_T:T_use]

        # 4. Group stats
        gs_dap = min(group_stats_dap, T_use - 1)
        group_stats = compute_group_stats_at_dap(fid_array, ndvi_concat, gs_dap)

        # 5. Run fine-tuning driver for this DAP window
        driver_kwargs = {k: v for k, v in ft_kwargs.items() if k != "rbf_params"}
        online_finetune_driver(
            model=model,
            ndvi_array=ndvi_concat,
            group_stats_array=group_stats,
            save_dir=save_dir,
            dap_start=current_dap,
            dap_end=min(current_dap + finetune_window, end_dap),
            dap_step=dap_step,
            max_dap=T_use,
            seed_dap=current_dap,
            device=device,
            **driver_kwargs,
        )

        current_dap += dap_step

    print(f"\n[DONE] Streaming fine-tuning complete. Results at: {save_dir}")
