# src/normet/analysis/decomposition.py
from __future__ import annotations

import os
import time
import numpy as np
import pandas as pd
from typing import Any, Dict, List, Optional, Tuple

from .normalise import normalise
from ..utils.features import extract_features
from ..model.train import build_model
from ..utils.prepare import process_date, add_date_variables
from ..utils.logging import get_logger

log = get_logger(__name__)


def _effective_cores(n_cores: Optional[int]) -> int:
    """Resolve parallel worker count (>=1)."""
    return max(1, n_cores if n_cores is not None else (os.cpu_count() or 2) - 1)

def decompose(
    method: str = "emission",
    df: Optional[pd.DataFrame] = None,
    model: Optional[object] = None,
    value: str = "value",
    feature_names: Optional[List[str]] = None,
    backend: Optional[str] = None,
    split_method: str = "random",
    fraction: float = 0.75,
    model_config: Optional[Dict[str, Any]] = None,
    n_samples: int = 300,
    seed: int = 7654321,
    importance_ascending: bool = False,
    n_cores: Optional[int] = None,
    memory_save: bool = False,
    verbose: bool = True,
) -> pd.DataFrame:
    """
    High-level wrapper for time series decomposition.

    Parameters
    ----------
    method : {"emission", "meteorology"}
        Decomposition strategy.
    df : pandas.DataFrame
        Input data with datetime and target column.
    model : object, optional
        Pre-trained model. If None, a new model will be trained.
    value : str
        Target column name.
    feature_names : List[str], optional
        Features used for training and decomposition.
    backend : str, optional
        Backend used for training if model is None.
    split_method, fraction, model_config, n_samples, seed, n_cores, memory_save, verbose
        Passed to underlying decomposition functions.

    Returns
    -------
    pandas.DataFrame
        Decomposed result.
    """
    # --- 1. Validate Inputs ---
    if df is None:
        raise ValueError("`df` must be provided.")
    if value is None:
        raise ValueError("`value` must be provided.")
    if model is None and feature_names is None:
        raise ValueError("Either `model` or `feature_names` must be provided.")
    if model is None and backend is None:
        raise ValueError("When training a model, `backend` must be specified.")

    # --- 2. Dispatch ---
    if method == "emission":
        return decom_emi(
            df=df, model=model, value=value, feature_names=feature_names,
            backend=backend, split_method=split_method, fraction=fraction,
            model_config=model_config, n_samples=n_samples, seed=seed,
            n_cores=n_cores, memory_save=memory_save, verbose=verbose
        )

    if method == "meteorology":
        return decom_met(
            df=df, model=model, value=value, feature_names=feature_names,
            backend=backend, split_method=split_method, fraction=fraction,
            model_config=model_config, n_samples=n_samples, seed=seed,
            importance_ascending=importance_ascending,
            n_cores=n_cores, memory_save=memory_save, verbose=verbose
        )

    # --- 3. Unsupported Method ---
    raise ValueError(
        f"Unsupported decomposition method: '{method}'. "
        "Must be one of 'emission' or 'meteorology'."
    )




def decom_emi(
    df: Optional[pd.DataFrame] = None,
    model: Optional[object] = None,
    value: str = "value",
    feature_names: Optional[List[str]] = None,
    backend: Optional[str] = None,
    split_method: str = "random",
    fraction: float = 0.75,
    model_config: Optional[Dict[str, Any]] = None,
    n_samples: int = 300,
    seed: int = 7654321,
    n_cores: Optional[int] = None,
    memory_save: bool = False,
    verbose: bool = True,
) -> pd.DataFrame:
    if df is None:
        raise ValueError("`df` must be provided.")
    if value is None:
        raise ValueError("`value` (target column name) must be provided.")
    if model is None and feature_names is None:
        raise ValueError("Either `model` or `feature_names` must be provided.")
    if model is None and backend is None:
        raise ValueError("When training a model, `backend` must be specified.")

    df_work = process_date(df.copy()) if "date" not in df.columns else df.copy()
    if "date" not in df_work.columns:
        raise ValueError("Could not find or create a 'date' column.")

    if value not in df_work.columns:
        raise ValueError(f"`df` does not contain the target column '{value}'.")

    observed_series = df_work[value].copy()
    if value != "value":
        df_work = df_work.rename(columns={value: "value"})

    mask_valid = df_work["date"].notna() & df_work["value"].notna()
    df_work = df_work.loc[mask_valid].sort_values("date").reset_index(drop=True)
    observed_series = observed_series.loc[mask_valid].reset_index(drop=True)

    if feature_names:
        missing_time_vars = [v for v in ["date_unix", "day_julian", "weekday", "hour"]
                             if v in feature_names and v not in df_work.columns]
        if missing_time_vars:
            try:
                df_work = add_date_variables(df_work)
                (log.info if verbose else log.debug)("Generated time variables: %s", missing_time_vars)
            except Exception:
                log.warning("Could not generate some time features: %s", missing_time_vars, exc_info=False)

    if model is None:
        (log.info if verbose else log.debug)(
            "Training model via backend='%s' with features=%d...", backend, len(feature_names)
        )
        df_work, model = build_model(
            df=df_work,
            value="value",
            backend=backend,
            feature_names=feature_names,
            split_method=split_method,
            fraction=fraction,
            model_config=model_config,
            seed=seed,
            n_cores=n_cores,
            verbose=verbose,
        )

    try:
        model_feats = [str(c) for c in extract_features(model)]
    except Exception:
        if not feature_names:
            raise ValueError("Cannot infer model features; please provide `feature_names`.")
        model_feats = [str(c) for c in feature_names]

    model_feats = [c for c in model_feats if c in df_work.columns]
    if not model_feats:
        raise ValueError("No valid model features found in the provided `df` for decomposition.")

    result = (
        pd.DataFrame({"date": df_work["date"].to_numpy(), "observed": observed_series.to_numpy()})
        .set_index("date")
        .sort_index()
    )

    time_vars_order = ["base", "date_unix", "day_julian", "weekday", "hour"]
    present_time_vars = ["base"] + [v for v in time_vars_order[1:] if v in model_feats and v in df_work.columns]

    n_cores_eff = _effective_cores(n_cores)
    start = time.time()

    for i, var_to_exclude in enumerate(present_time_vars, start=1):
        resample_vars = [v for v in model_feats if v != var_to_exclude and v != "value"]

        elapsed = time.time() - start
        eta = (elapsed / max(i - 1, 1)) * (len(present_time_vars) - (i - 1)) if i > 1 else None
        eta_str = "" if eta is None else (
            f" | ETA: {eta:.1f}s" if eta < 60 else f" | ETA: {eta/60:.1f}m" if eta < 3600 else f" | ETA: {eta/3600:.1f}h"
        )
        (log.info if verbose else log.debug)(
            "%s: Decomposing %s%s", pd.Timestamp.now().strftime("%Y-%m-%d %H:%M:%S"), var_to_exclude, eta_str
        )

        df_norm = normalise(
            df=df_work,
            model=model,
            feature_names=model_feats,
            variables_resample=resample_vars,
            n_samples=n_samples,
            replace=True,
            aggregate=True,
            seed=seed,
            n_cores=n_cores_eff,
            weather_df=None,
            memory_save=memory_save,
            verbose=False,
        )
        if "normalised" not in df_norm.columns:
            log.exception("`normalise` did not return 'normalised' column (aggregate=True).")
            raise RuntimeError("`normalise` must return a DataFrame with column 'normalised' when aggregate=True.")

        result[var_to_exclude] = df_norm.reindex(result.index)["normalised"].to_numpy()

    result["emi_total"] = result.get("hour", result["observed"])

    for a, b, out in [
        ("hour", "weekday", "hour"),
        ("weekday", "day_julian", "weekday"),
        ("day_julian", "date_unix", "day_julian"),
        ("date_unix", "base", "date_unix"),
    ]:
        if a in result.columns and b in result.columns:
            result[out] = result[a] - result[b]

    base_mean = float(result["base"].mean())
    result["emi_noise"] = result["base"] - base_mean
    result["emi_base"] = base_mean
    del result["base"]

    return result



def decom_met(
    df: Optional[pd.DataFrame] = None,
    model: Optional[object] = None,
    value: str = "value",
    feature_names: Optional[List[str]] = None,
    backend: Optional[str] = None,
    split_method: str = "random",
    fraction: float = 0.75,
    model_config: Optional[Dict[str, Any]] = None,
    n_samples: int = 300,
    seed: int = 7654321,
    importance_ascending: bool = False,
    n_cores: Optional[int] = None,
    memory_save: bool = False,
    verbose: bool = True,
) -> pd.DataFrame:
    if df is None:
        raise ValueError("`df` must be provided.")
    if value is None:
        raise ValueError("`value` (target column name) must be provided.")
    if model is None and feature_names is None:
        raise ValueError("Either `model` or `feature_names` must be provided.")
    if model is None and backend is None:
        raise ValueError("When training a model, `backend` must be specified.")

    df = df.copy()
    if "date" not in df.columns:
        df = process_date(df)
    df = df[df["date"].notna()].sort_values("date").reset_index(drop=True)

    if value not in df.columns:
        raise ValueError(f"`df` does not contain the target column '{value}'.")
    observed_series = df[value].copy()

    df_work = df.copy()
    if value != "value":
        df_work = df_work.rename(columns={value: "value"})

    # Add missing time variables if needed
    if feature_names:
        time_vars = ["date_unix", "day_julian", "weekday", "hour"]
        missing_time_vars = [v for v in time_vars if v in feature_names and v not in df_work.columns]
        if missing_time_vars:
            try:
                df_work = add_date_variables(df_work)
                (log.info if verbose else log.debug)("Generated time variables: %s", missing_time_vars)
            except Exception:
                log.warning("Missing time features not generated: %s", missing_time_vars, exc_info=False)

    # Train model if needed
    if model is None:
        (log.info if verbose else log.debug)(
            "Training model via backend='%s' with features=%d...", backend, len(feature_names)
        )
        df_work, model = build_model(
            df=df_work,
            value="value",
            backend=backend,
            feature_names=feature_names,
            split_method=split_method,
            fraction=fraction,
            model_config=model_config,
            seed=seed,
            verbose=verbose,
        )

    # Extract feature importance order
    try:
        feat_sorted = extract_features(model, importance_ascending=importance_ascending)
    except Exception:
        if not feature_names:
            raise ValueError("Cannot infer model features; please provide `feature_names`.")
        feat_sorted = list(feature_names)

    feat_sorted = [f for f in feat_sorted if f in df_work.columns]
    if not feat_sorted:
        raise ValueError("No valid model features found in `df`.")

    # Exclude time variables from meteorological contribution set
    time_vars = {"hour", "weekday", "day_julian", "date_unix"}
    contrib_candidates = [f for f in feat_sorted if f not in time_vars]

    result = (
        pd.DataFrame({"date": df_work["date"].to_numpy(), "observed": observed_series.to_numpy()})
        .set_index("date")
        .sort_index()
    )

    n_cores_eff = _effective_cores(n_cores)
    decomp_order = ["emi_total"] + contrib_candidates[:]
    resample_vars = contrib_candidates[:]

    start = time.time()
    tmp: Dict[str, np.ndarray] = {}

    for i, var in enumerate(decomp_order, start=1):
        if var != "emi_total" and var in resample_vars:
            resample_vars = [v for v in resample_vars if v != var]

        elapsed = time.time() - start
        eta = (elapsed / max(i - 1, 1)) * (len(decomp_order) - (i - 1)) if i > 1 else None
        eta_str = "" if eta is None else (
            f" | ETA: {eta:.1f}s" if eta < 60 else f" | ETA: {eta/60:.1f}m" if eta < 3600 else f" | ETA: {eta/3600:.1f}h"
        )
        (log.info if verbose else log.debug)(
            "%s: Decomposing %s%s", pd.Timestamp.now().strftime("%Y-%m-%d %H:%M:%S"), var, eta_str
        )

        df_norm = normalise(
            df=df_work,
            model=model,
            feature_names=feat_sorted,
            variables_resample=resample_vars,
            n_samples=n_samples,
            replace=True,
            aggregate=True,
            seed=seed,
            n_cores=n_cores_eff,
            weather_df=None,
            memory_save=memory_save,
            verbose=False,
        )
        if "normalised" not in df_norm.columns:
            log.exception("`normalise` did not return 'normalised' column (aggregate=True).")
            raise RuntimeError("`normalise` must return a DataFrame with column 'normalised' when aggregate=True.")

        tmp[var] = df_norm.reindex(result.index)["normalised"].to_numpy()

    result["emi_total"] = tmp["emi_total"]

    prev_key = "emi_total"
    for feat in contrib_candidates:
        result[feat] = tmp[feat] - tmp[prev_key]
        prev_key = feat

    result["met_total"] = result["observed"] - result["emi_total"]
    result["met_base"] = float(result["met_total"].mean())
    contrib_sum = result[contrib_candidates].sum(axis=1) if contrib_candidates else 0.0
    result["met_noise"] = result["met_total"] - (result["met_base"] + contrib_sum)

    return result
