# main.py
import json, traceback
import pandas as pd
import numpy as np
import io, tempfile, textwrap
import statsmodels.api as sm

from fastapi import FastAPI, HTTPException, Body, UploadFile, File, Form, Response, status, Request
from fastapi.responses import FileResponse, PlainTextResponse, ORJSONResponse
from fastapi.concurrency import run_in_threadpool
from datetime import datetime
from typing import List, Dict, Any, Optional
from dataclasses import dataclass

# Predefined globals
stored_model = None
stored_result_join = None

@dataclass
class DemoFitResult:
    formula: str
    alpha: List[List[float]]
    beta: List[List[float]]
    headers_alpha: List[str]
    headers_beta: List[str]
    Performance: List[List[float]]
    R2_individual: List[float]
    R2_individual_labels: List[str]

@dataclass
class DemoFitJoin:
    result: DemoFitResult

class DATFIDModel:
    """
    Demo-safe replacement for the private DATFID model.
    Uses simple OLS with optional lagged features and trend.
    """
    def __init__(
        self,
        df: pd.DataFrame,
        id_col: str,
        time_col: str,
        y: str,
        lag_y: Any = None,
        lagged_features: Any = None,
        current_features: Any = None,
        filter_by_significance: bool = False,
        meanvar_test: bool = False,
        signif: float = 0.05,
    ):
        self.df = df.copy()
        self.id_col = id_col
        self.time_col = time_col
        self.y = y
        self.lag_y = lag_y
        self.lagged_features = lagged_features if isinstance(lagged_features, dict) else {}
        if current_features == "all":
            self.current_features = "all"
        elif isinstance(current_features, list):
            self.current_features = current_features
        else:
            self.current_features = []
        self.filter_by_significance = filter_by_significance
        self.meanvar_test = meanvar_test
        self.signif = signif
        self._fitted_model = None
        self._train_columns: List[str] = []
        self._last_y_by_id: Dict[str, float] = {}
        self._last_y_global: float = 0.0

    def _get_lag_int(self, value: Any, default: int = 1) -> int:
        try:
            out = int(value)
            return out if out > 0 else default
        except Exception:
            return default

    def _sorted(self, df: pd.DataFrame) -> pd.DataFrame:
        if self.id_col in df.columns and self.time_col in df.columns:
            return df.sort_values([self.id_col, self.time_col]).copy()
        if self.time_col in df.columns:
            return df.sort_values([self.time_col]).copy()
        return df.copy()

    def _trend(self, df: pd.DataFrame) -> pd.Series:
        if self.time_col in df.columns:
            ts = pd.to_datetime(df[self.time_col], errors="coerce")
            if ts.notna().any():
                base = ts.min()
                return (ts - base).dt.days.fillna(0.0).astype(float)
        return pd.Series(np.arange(len(df), dtype=float), index=df.index)

    def _resolve_current_features(self, df: pd.DataFrame) -> List[str]:
        if self.current_features == "all":
            return [
                c for c in df.columns
                if c not in {self.id_col, self.time_col, self.y}
                and pd.api.types.is_numeric_dtype(df[c])
            ]
        return [c for c in self.current_features if c in df.columns]

    def _build_matrix(self, df: pd.DataFrame, is_forecast: bool) -> pd.DataFrame:
        work = self._sorted(df)
        x = pd.DataFrame(index=work.index)
        x["trend_index"] = self._trend(work)

        for col in self._resolve_current_features(work):
            x[col] = pd.to_numeric(work[col], errors="coerce")

        if self.y in work.columns:
            lag_y_int = self._get_lag_int(self.lag_y, default=1) if self.lag_y else None
            if lag_y_int:
                lag_name = f"{self.y}_lag_{lag_y_int}"
                if self.id_col in work.columns:
                    x[lag_name] = work.groupby(self.id_col, sort=False)[self.y].shift(lag_y_int)
                else:
                    x[lag_name] = work[self.y].shift(lag_y_int)

        for feat, lag in (self.lagged_features or {}).items():
            if feat not in work.columns:
                continue
            lag_int = self._get_lag_int(lag, default=1)
            col_name = f"{feat}_lag_{lag_int}"
            if self.id_col in work.columns:
                x[col_name] = work.groupby(self.id_col, sort=False)[feat].shift(lag_int)
            else:
                x[col_name] = work[feat].shift(lag_int)

        x = x.apply(pd.to_numeric, errors="coerce")
        if is_forecast:
            x = x.fillna(0.0)
        return x

    def fit(self) -> DemoFitJoin:
        if self.y not in self.df.columns:
            raise ValueError(f"Target column '{self.y}' is missing.")

        train = self._sorted(self.df)
        x = self._build_matrix(train, is_forecast=False)
        y = pd.to_numeric(train[self.y], errors="coerce")
        valid = y.notna()
        if x.shape[1] > 0:
            valid = valid & x.notna().all(axis=1)
        x = x.loc[valid].copy()
        y = y.loc[valid].copy()

        if len(y) < 3:
            raise ValueError("Not enough valid rows to fit demo model (need >= 3).")

        x_const = sm.add_constant(x, has_constant="add")
        self._fitted_model = sm.OLS(y.astype(float), x_const.astype(float)).fit()
        self._train_columns = list(x_const.columns)

        if self.id_col in train.columns:
            last_vals = train.groupby(self.id_col, sort=False)[self.y].last().dropna()
            self._last_y_by_id = {str(k): float(v) for k, v in last_vals.items()}
        self._last_y_global = float(y.iloc[-1]) if len(y) else 0.0

        params = self._fitted_model.params
        bse = self._fitted_model.bse
        tvals = self._fitted_model.tvalues
        pvals = self._fitted_model.pvalues

        const_name = "const" if "const" in params.index else params.index[0]
        beta_names = [n for n in params.index if n != const_name]

        alpha = [[float(params.get(const_name, 0.0))], [float(bse.get(const_name, 0.0))], [float(tvals.get(const_name, 0.0))], [float(pvals.get(const_name, 1.0))]]
        beta = [
            [float(params.get(n, 0.0)) for n in beta_names],
            [float(bse.get(n, 0.0)) for n in beta_names],
            [float(tvals.get(n, 0.0)) for n in beta_names],
            [float(pvals.get(n, 1.0)) for n in beta_names],
        ]

        pred = self._fitted_model.predict(x_const)
        mse = float(np.mean((y - pred) ** 2))
        mae = float(np.mean(np.abs(y - pred)))
        r2 = float(getattr(self._fitted_model, "rsquared", 0.0))
        r2_adj = float(getattr(self._fitted_model, "rsquared_adj", r2))
        perf = [
            [r2, r2],
            [r2_adj, r2_adj],
            [r2, r2_adj],
            [mse, mse],
            [mae, mae],
        ]

        r2_individual: List[float] = []
        r2_labels: List[str] = []
        if self.id_col in train.columns:
            joined = pd.DataFrame({
                self.id_col: train.loc[valid, self.id_col].astype(str),
                "_y": y.values,
                "_p": pred.values,
            })
            for id_val, sub in joined.groupby(self.id_col, sort=False):
                den = float(((sub["_y"] - sub["_y"].mean()) ** 2).sum())
                if den <= 0:
                    r2_i = 0.0
                else:
                    num = float(((sub["_y"] - sub["_p"]) ** 2).sum())
                    r2_i = 1.0 - (num / den)
                r2_labels.append(str(id_val))
                r2_individual.append(r2_i)

        formula = f"{self.y} ~ " + " + ".join(self._train_columns)
        fit_result = DemoFitResult(
            formula=formula,
            alpha=alpha,
            beta=beta,
            headers_alpha=[const_name],
            headers_beta=beta_names,
            Performance=perf,
            R2_individual=r2_individual,
            R2_individual_labels=r2_labels,
        )
        return DemoFitJoin(result=fit_result)

    def forecast(self, extern_self: Any, df_forecast: pd.DataFrame) -> pd.DataFrame:
        if self._fitted_model is None:
            raise ValueError("Model not fitted.")

        out = self._sorted(df_forecast).copy()
        x = self._build_matrix(out, is_forecast=True)
        x_const = sm.add_constant(x, has_constant="add")

        for col in self._train_columns:
            if col not in x_const.columns:
                x_const[col] = 0.0
        x_const = x_const[self._train_columns].astype(float)

        pred = self._fitted_model.predict(x_const)
        out[f"{self.y}_forecast"] = np.asarray(pred, dtype=float)
        out["forecast"] = out[f"{self.y}_forecast"]
        return out

def _maybe_json_list(s: Optional[str]):
    if s is None or s == "":
        return []
    s = s.strip()
    if s.lower() == "all":
        return "all"
    try:
        val = json.loads(s)
        if isinstance(val, list):
            return val
        return []
    except Exception:
        # allow comma-separated as a fallback
        return [x.strip() for x in s.split(",") if x.strip()]
    
def _maybe_json_dict(s: Optional[str]):
    if not s:
        return {}
    try:
        val = json.loads(s)
        return val if isinstance(val, dict) else {}
    except Exception:
        return {}
    
def _read_table_from_upload(upload: UploadFile) -> pd.DataFrame:
    name = (upload.filename or "").lower()
    data = upload.file.read()
    bio = io.BytesIO(data)
    if name.endswith(".csv"):
        return pd.read_csv(bio)
    # default to Excel (supports .xls, .xlsx)
    return pd.read_excel(bio)

def _result_to_text(result_obj: Any) -> str:
    """
    Build a compact, readable DATFID model summary.

    Expected fields in `result_obj` (object with attributes or a dict):
      - formula: str
      - alpha: 2D array-like (rows ~ [Estimate, SE, T, P], columns = time-invariant features)
      - beta:  2D array-like (rows ~ [Estimate, SE, T, P], columns = time-variant features)
      - headers_alpha: list[str] (names for alpha columns)
      - headers_beta:  list[str] (names for beta columns)
      - Performance: 2D array-like with the last 5 rows (in this order):
            R2 within, R2 between, R2 overall, MSE, MAE
        and 2 columns:
            2SFE, 2SFE_c
      - R2_individual: 1D array-like of per-individual R² (optional)
      - R2_individual_labels: list[str] of same length as R2_individual (optional)
        If provided, labels will be shown instead of ID numbers in the summary.

    The function tolerates:
      - dict or object input
      - missing headers (falls back to generic names)
      - extra rows in alpha/beta (truncated to 4)
      - extra rows in Performance (only last 5 kept)
    """

    # --- Safe getters for both dicts and objects
    def get(key, default=None):
        if isinstance(result_obj, dict):
            return result_obj.get(key, default)
        return getattr(result_obj, key, default)

    def as_2d(a):
        if a is None:
            return np.empty((0, 0), dtype=float)
        arr = np.asarray(a, dtype=float)
        if arr.ndim == 1:
            arr = arr[None, :]
        return arr

    # --- Pull fields
    formula = (get("formula", "") or "").strip()

    alpha_arr = as_2d(get("alpha"))
    beta_arr  = as_2d(get("beta"))
    perf_arr  = as_2d(get("Performance"))

    headers_alpha = list(get("headers_alpha", [])) or [f"Alpha_{i+1}" for i in range(alpha_arr.shape[1])]
    headers_beta  = list(get("headers_beta", []))  or [f"Beta_{i+1}"  for i in range(beta_arr.shape[1])]

    r2_individual = get("R2_individual", None)
    r2_labels     = get("R2_individual_labels", None)

    # --- Shape/label guards
    row_labels = ["Estimate", "Standard Error", "T statistic", "P value"]

    if alpha_arr.shape[0] > 4:
        alpha_arr = alpha_arr[:4, :]
    if beta_arr.shape[0] > 4:
        beta_arr = beta_arr[:4, :]

    # If perf has >5 rows, keep the last 5 (assumes metrics are at the end)
    if perf_arr.shape[0] >= 5:
        perf_arr = perf_arr[-5:, :]
    perf_rows = ["R2 within", "R2 between", "R2 overall", "MSE", "MAE"]
    perf_cols = ["2SFE", "2SFE_c"]
    # Guard columns
    n_perf_cols = min(perf_arr.shape[1], 2)
    perf_cols = perf_cols[:n_perf_cols]

    # --- Build tables
    alpha_df = pd.DataFrame(alpha_arr, index=row_labels[:alpha_arr.shape[0]],
                            columns=headers_alpha[:alpha_arr.shape[1]])
    beta_df  = pd.DataFrame(beta_arr,  index=row_labels[:beta_arr.shape[0]],
                            columns=headers_beta[:beta_arr.shape[1]])
    perf_df  = pd.DataFrame(perf_arr,  index=perf_rows[:perf_arr.shape[0]],
                            columns=perf_cols)

    # --- R² summary (min/median/max)
    r2_lines = []
    if r2_individual is not None:
        r2_vals = np.asarray(r2_individual, dtype=float).ravel()
        if r2_vals.size > 0 and np.isfinite(r2_vals).any():
            # Prepare labels (either provided or fallback to 1-based IDs)
            if r2_labels and len(r2_labels) == r2_vals.size:
                ids = np.array(list(r2_labels), dtype=object)
                def _lab(idx): return str(ids[idx])
            else:
                def _lab(idx): return f"ID {idx+1}"

            r2_min = float(np.nanmin(r2_vals))
            r2_med = float(np.nanmedian(r2_vals))
            r2_max = float(np.nanmax(r2_vals))

            # Nearest index to each statistic (handles non-exact medians)
            imin = int(np.nanargmin(r2_vals))
            imed = int(np.nanargmin(np.abs(r2_vals - r2_med)))
            imax = int(np.nanargmax(r2_vals))

            r2_lines = [
                f"min ({_lab(imin)}): {r2_min:.6g}",
                f"median ({_lab(imed)}): {r2_med:.6g}",
                f"max ({_lab(imax)}): {r2_max:.6g}",
            ]

    # --- Pretty printers
    ffmt = lambda x: f"{x:.6g}"
    def _df_text(df: pd.DataFrame) -> str:
        # Right-justified columns; scientific format where appropriate
        return df.to_string(justify="right", float_format=ffmt)

    # --- Compose report
    parts = []
    parts.append("DATFID Fit Result")
    parts.append(f"Generated: {datetime.utcnow().isoformat()}Z")
    parts.append("=" * 72)
    parts.append("=== Model Summary ===\n")
    parts.append("Formula:")
    parts.append(f" {formula}\n")

    parts.append("Alpha (time invariant):")
    parts.append(_df_text(alpha_df) + "\n")

    parts.append("Beta (time variant):")
    parts.append(_df_text(beta_df) + "\n")

    parts.append("Performance metrics:")
    parts.append(_df_text(perf_df) + "\n")

    if r2_lines:
        parts.append("Individual R² summary:")
        parts.extend(r2_lines)

    return ("\n".join(parts)).rstrip() + "\n"

# ---------- FastAPI app ----------
app = FastAPI(
    title="DATFID API",
    description="Public demo API",
    docs_url="/docs",
    redoc_url=None,
    default_response_class=ORJSONResponse,
)

@app.get("/")
def root():
    return {"message": "DATFID API is alive."}

# ✅ Validate directly here
@app.get("/secure-ping/")
def secure_ping():
    return {"ok": True}

# ---------- Model endpoints (guarded) ----------
@app.post("/modelfit/")
async def modelfit(
    df: List[Dict] = Body(...),
    id_col: str = Body(...),
    time_col: str = Body(...),
    y: str = Body(...),
    lag_y: Any = Body(None),
    lagged_features: Any = Body({}),
    current_features: Any = Body([]),
    filter_by_significance: bool = Body(False),
    meanvar_test: bool = Body(False),
    signif: Any = Body(0.05),
):
    global stored_model, stored_result_join
    try:
        sig_val = float(signif)
    except (TypeError, ValueError):
        sig_val = 0.05
    df1 = pd.DataFrame(df)
    try:
        model = DATFIDModel(
            df=df1,
            id_col=id_col,
            time_col=time_col,
            y=y,
            lag_y=lag_y,
            lagged_features=lagged_features,
            current_features=current_features,
            filter_by_significance=filter_by_significance,
            meanvar_test=meanvar_test,
            signif=sig_val,
        )
        result_join = model.fit()
        result = result_join.result

        # save for forecast
        stored_model = model
        stored_result_join = result_join

        # jsonify result object
        result_dict = {}
        for k, v in result.__dict__.items():
            if isinstance(v, pd.DataFrame):
                result_dict[k] = v.to_dict(orient="records")
            elif isinstance(v, pd.Series):
                result_dict[k] = v.to_dict()
            elif isinstance(v, (list, dict, str, int, float, bool, type(None))):
                result_dict[k] = v
            else:
                result_dict[k] = str(v)
        return result_dict
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error during model fit: {str(e)}")

@app.post("/modelforecast/")
async def modelforecast(
    payload: Any = Body(...),
):
    global stored_model, stored_result_join
    if stored_model is None or stored_result_join is None:
        raise HTTPException(status_code=400, detail="Model not fitted. Call /modelfit/ first.")
    try:
        # Accept both payload styles:
        # 1) raw list: [...]
        # 2) wrapped dict: {"df_forecast": [...]}
        if isinstance(payload, dict) and "df_forecast" in payload:
            df_forecast = payload.get("df_forecast")
        else:
            df_forecast = payload
        if not isinstance(df_forecast, list):
            raise HTTPException(
                status_code=422,
                detail="Payload must be a list of rows or {'df_forecast': [rows]}",
            )
        df_forecast1 = pd.DataFrame(df_forecast)
        forecast_df = stored_model.forecast(extern_self=stored_result_join, df_forecast=df_forecast1)
        return forecast_df.to_dict(orient="records")
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error during model forecast: {str(e)}")

@app.post("/modelfit-file/")
async def modelfit_file(
    file: UploadFile = File(...),
    id_col: str = Form(...),
    time_col: str = Form(...),
    y: str = Form(...),
    lag_y: str = Form(""),
    lagged_features: str = Form(""),
    current_features: str = Form(""),
    filter_by_significance: str = Form("false"),
    meanvar_test: str = Form("false"),
    signif: str = Form("0.05"),
):
    global stored_model, stored_result_join
    try:
        df = _read_table_from_upload(file)

        # harmonize datetimes → str (like SDK)
        for col in df.columns:
            if pd.api.types.is_datetime64_any_dtype(df[col]):
                df[col] = df[col].astype(str)

        lagged = _maybe_json_dict(lagged_features)
        curr = _maybe_json_list(current_features)
        filt_sig = str(filter_by_significance).strip().lower() == "true"
        mv_test = str(meanvar_test).strip().lower() == "true"
        ly = None if (lag_y is None or lag_y.strip() == "") else lag_y.strip()
        try:
            sig_val = float(signif)
        except (TypeError, ValueError):
            sig_val = 0.05


        model = DATFIDModel(
            df=df,
            id_col=id_col,
            time_col=time_col,
            y=y,
            lag_y=ly,
            lagged_features=lagged,
            current_features=curr,
            filter_by_significance=filt_sig,
            meanvar_test=mv_test,
            signif=sig_val,
        )
        result_join = model.fit()
        stored_model = model
        stored_result_join = result_join

        # Build textual report
        report_text = _result_to_text(result_join.result)
        # Write to a temp file and return as attachment
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".txt")
        tmp.write(report_text.encode("utf-8"))
        tmp.flush(); tmp.close()

        fname = "result.txt"
        return FileResponse(
            tmp.name,
            media_type="text/plain; charset=utf-8",
            filename=fname,
        )
    except HTTPException:
        raise
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error during model fit (file): {str(e)}")
    
@app.post("/modelforecast-file/")
async def modelforecast_file(
    df_forecast: UploadFile = File(...),
):
    global stored_model, stored_result_join
    if stored_model is None or stored_result_join is None:
        raise HTTPException(status_code=400, detail="Model not fitted. Call /modelfit-file/ (or /modelfit/) first.")
    try:
        df_fc = _read_table_from_upload(df_forecast)

        # harmonize datetimes → str (like SDK)
        for col in df_fc.columns:
            if pd.api.types.is_datetime64_any_dtype(df_fc[col]):
                df_fc[col] = df_fc[col].astype(str)

        forecast_df = stored_model.forecast(extern_self=stored_result_join, df_forecast=df_fc)

        # Save to CSV and return (Excel-friendly)
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
        tmp_path = tmp.name
        tmp.close()  # we'll reopen with encoding

        # Build CSV in memory
        buf = io.StringIO()
        forecast_df.to_csv(buf, index=False)  # keep commas
        csv_text = buf.getvalue()

        # Excel friendliness:
        #  1) 'sep=,' header makes Excel use commas even in ; locales
        #  2) Also prevents 'ID' being the first two characters -> avoids SYLK warning
        csv_text = "sep=,\n" + csv_text

        # Write with BOM so Excel recognizes UTF-8
        with open(tmp_path, "w", encoding="utf-8-sig", newline="") as f:
            f.write(csv_text)

        return FileResponse(
            tmp_path,
            media_type="text/csv; charset=utf-8",
            filename="forecast.csv",
        )

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error during model forecast (file): {str(e)}")