"""Evaluate saved TFT checkpoints on one immutable validation/test snapshot.

This is a diagnostic tool: checkpoint ranking is computed from validation
metrics only.  Test metrics are printed separately to measure whether the
validation rule transfers; they never participate in the ranking.
"""

from __future__ import annotations

import argparse
import copy
import json
from dataclasses import replace
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from app.quality_gate import evaluate_quality_gate_metrics
from deep_learning.config import get_tft_config
from deep_learning.data.dataset import build_datasets, create_dataloaders
from deep_learning.models.tft_copper import load_tft_model
from deep_learning.training.direction_model import (
    apply_weekly_direction_model,
    predict_weekly_direction,
)
from deep_learning.training.metrics import (
    apply_weekly_sign_correction_np,
    cumulative_horizon,
    directional_accuracy,
    fit_direction_sign_calibration,
    fit_weekly_interval_scale,
)
from deep_learning.training.trainer import (
    WEEKLY_INTERVAL_CONDITIONING_FEATURE,
    _build_uniform_checkpoint_soup,
    _compute_test_metrics_from_quantiles,
    _validate_quantile_prediction_shape,
    _weekly_interval_conditioning_values,
)


def _actual_path(dataloader) -> np.ndarray:
    parts = [
        batch[1][0] if isinstance(batch[1], (list, tuple)) else batch[1]
        for batch in dataloader
    ]
    return torch.cat(parts).cpu().numpy()


def _predict_quantiles(model, dataloader, cfg) -> np.ndarray:
    prediction = model.predict(
        dataloader,
        mode="quantiles",
        trainer_kwargs={"logger": False, "enable_checkpointing": False},
    )
    prediction_np = prediction.cpu().numpy()
    _validate_quantile_prediction_shape(prediction_np, cfg)
    return prediction_np


def _prepare_candidate(
    *,
    prediction: np.ndarray,
    actual: np.ndarray,
    weekly_direction_model: dict,
    direction_probability: np.ndarray,
    conditioning_values: np.ndarray | None,
    cfg,
) -> tuple[np.ndarray, dict, dict, dict, dict]:
    direction_calibration = fit_direction_sign_calibration(
        actual,
        prediction,
        horizon=cfg.forecast.primary_horizon_days,
    )
    oriented = prediction * int(
        direction_calibration.get("direction_sign_multiplier", 1)
    )
    oriented = apply_weekly_sign_correction_np(
        oriented,
        float(direction_calibration.get("weekly_sign_threshold", 0.0)),
        horizon=cfg.forecast.primary_horizon_days,
    )

    selected_direction_model = copy.deepcopy(weekly_direction_model)
    selected_direction_model["enabled"] = False
    selected_direction_model["reason"] = "validation_selection_rejected"
    if selected_direction_model.get("coef") and len(direction_probability) == len(
        oriented
    ):
        candidate = apply_weekly_direction_model(
            oriented,
            direction_probability,
            threshold=float(
                selected_direction_model.get("decision_threshold", 0.50)
            ),
            horizon=cfg.forecast.primary_horizon_days,
        )
        actual_weekly = cumulative_horizon(
            actual,
            horizon=cfg.forecast.primary_horizon_days,
        )
        median_idx = len(cfg.model.quantiles) // 2
        base_da = directional_accuracy(
            actual_weekly,
            cumulative_horizon(
                oriented[:, :, median_idx],
                horizon=cfg.forecast.primary_horizon_days,
            ),
        )
        candidate_da = directional_accuracy(
            actual_weekly,
            cumulative_horizon(
                candidate[:, :, median_idx],
                horizon=cfg.forecast.primary_horizon_days,
            ),
        )
        candidate_rate = float(np.mean(direction_probability >= 0.50))
        selected_direction_model.update(
            {
                "validation_base_weekly_da": float(base_da),
                "validation_weekly_da": float(candidate_da),
                "validation_pred_positive_rate": candidate_rate,
            }
        )
        if (
            candidate_da >= 0.51
            and candidate_da >= base_da + 0.01
            and 0.25 <= candidate_rate <= 0.75
        ):
            selected_direction_model["enabled"] = True
            selected_direction_model["reason"] = "validation_improved"
            oriented = candidate

    interval_calibration = fit_weekly_interval_scale(
        actual,
        oriented,
        quantiles=tuple(cfg.model.quantiles),
        horizon=cfg.forecast.primary_horizon_days,
        weekly_median_cap=cfg.weekly_loss.weekly_median_cap,
        conditioning_values=conditioning_values,
        conditioning_feature=(
            WEEKLY_INTERVAL_CONDITIONING_FEATURE
            if conditioning_values is not None
            else None
        ),
    )
    metrics = _compute_test_metrics_from_quantiles(
        actual,
        oriented,
        cfg,
        weekly_interval_scale=float(
            interval_calibration.get("weekly_interval_scale", 1.0)
        ),
        weekly_interval_calibration=interval_calibration,
        weekly_interval_conditioning_values=conditioning_values,
    )
    return (
        oriented,
        metrics,
        direction_calibration,
        interval_calibration,
        selected_direction_model,
    )


def _validation_rank(metrics: dict, checkpoint_loss: float) -> tuple:
    """Rank without test labels: fewer contract misses, then safer margins."""
    _passed, reasons = evaluate_quality_gate_metrics(metrics)

    # Proper interval score and bounded scale are continuous tie-breakers.
    return (
        len(reasons),
        float(metrics["weekly_pi80_interval_score"]),
        float(metrics["weekly_mae_vs_naive_zero"]),
        -float(metrics["weekly_sharpe_ratio"]),
        float(checkpoint_loss),
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("artifact_dir", type=Path)
    parser.add_argument("--output", type=Path, help="Write JSON results to this path")
    parser.add_argument("--include-soups", action="store_true")
    parser.add_argument("--candidate", help="Evaluate only one checkpoint filename")
    parser.add_argument(
        "--weekly-median-cap-factor",
        type=float,
        default=1.0,
        help="Diagnostic multiplier applied to the recorded weekly median cap",
    )
    args = parser.parse_args()

    artifact_dir = args.artifact_dir.resolve()
    snapshot_path = artifact_dir / "feature_snapshot.pkl"
    snapshot_meta = json.loads(
        snapshot_path.with_suffix(".pkl.json").read_text(encoding="utf-8")
    )
    run_meta = json.loads(
        (artifact_dir / "tft_metadata.json").read_text(encoding="utf-8")
    )

    frame = pd.read_pickle(snapshot_path)
    cfg = get_tft_config()
    effective_weekly_cap = (
        float(run_meta["config"]["weekly_median_cap"])
        * args.weekly_median_cap_factor
    )
    if effective_weekly_cap <= 0.0:
        raise ValueError("The effective weekly median cap must be positive")
    cfg = replace(
        cfg,
        weekly_loss=replace(
            cfg.weekly_loss,
            weekly_median_cap=effective_weekly_cap,
        ),
    )
    training_ds, validation_ds, test_ds = build_datasets(
        frame,
        snapshot_meta["time_varying_unknown_reals"],
        snapshot_meta["time_varying_known_reals"],
        snapshot_meta["target_cols"],
        cfg,
    )
    _, val_dl, test_dl = create_dataloaders(
        training_ds,
        validation_ds,
        test_ds,
        cfg,
    )
    val_actual = _actual_path(val_dl)
    test_actual = _actual_path(test_dl)

    n_rows = len(frame)
    test_size = int(n_rows * cfg.training.test_ratio)
    val_size = int(n_rows * cfg.training.val_ratio)
    train_size = n_rows - val_size - test_size
    train_cutoff = int(frame["time_idx"].iloc[train_size - 1])
    val_cutoff = int(frame["time_idx"].iloc[train_size + val_size - 1])
    max_time_idx = int(frame["time_idx"].iloc[-1])

    base_direction_model = run_meta["weekly_direction_model"]
    val_direction_probability = predict_weekly_direction(
        base_direction_model,
        frame,
        start_exclusive=train_cutoff - 1,
        end_inclusive=val_cutoff - cfg.forecast.primary_horizon_days,
    )
    test_direction_probability = predict_weekly_direction(
        base_direction_model,
        frame,
        start_exclusive=val_cutoff - 1,
        end_inclusive=max_time_idx - cfg.forecast.primary_horizon_days,
    )
    val_conditioning = _weekly_interval_conditioning_values(
        frame,
        feature_name=WEEKLY_INTERVAL_CONDITIONING_FEATURE,
        start_exclusive=train_cutoff - 1,
        end_inclusive=val_cutoff - cfg.forecast.primary_horizon_days,
        expected_count=len(val_actual),
    )
    test_conditioning = _weekly_interval_conditioning_values(
        frame,
        feature_name=WEEKLY_INTERVAL_CONDITIONING_FEATURE,
        start_exclusive=val_cutoff - 1,
        end_inclusive=max_time_idx - cfg.forecast.primary_horizon_days,
        expected_count=len(test_actual),
    )

    checkpoint_specs = []
    checkpoints = sorted(
        (artifact_dir / "checkpoints").glob("tft-asro-*.ckpt"),
        key=lambda path: float(path.stem.rsplit("=", 1)[-1]),
    )
    for checkpoint in checkpoints:
        checkpoint_specs.append(
            (checkpoint, float(checkpoint.stem.rsplit("=", 1)[-1]))
        )
    if args.include_soups:
        for count in range(2, len(checkpoints) + 1):
            source = checkpoints[:count]
            soup_path = (
                artifact_dir
                / "checkpoints"
                / f"validation-top-{count}-soup.ckpt"
            )
            _build_uniform_checkpoint_soup(source, soup_path)
            checkpoint_specs.append(
                (
                    soup_path,
                    float(
                        np.mean(
                            [
                                float(path.stem.rsplit("=", 1)[-1])
                                for path in source
                            ]
                        )
                    ),
                )
            )

    if args.candidate:
        checkpoint_specs = [
            spec for spec in checkpoint_specs if spec[0].name == args.candidate
        ]
        if not checkpoint_specs:
            raise FileNotFoundError(
                f"Checkpoint candidate not found: {args.candidate}"
            )

    rows = []
    for checkpoint, loss in checkpoint_specs:
        model = load_tft_model(str(checkpoint))
        val_prediction = _predict_quantiles(model, val_dl, cfg)
        (
            _,
            val_metrics,
            direction_calibration,
            interval_calibration,
            selected_direction_model,
        ) = _prepare_candidate(
            prediction=val_prediction,
            actual=val_actual,
            weekly_direction_model=base_direction_model,
            direction_probability=val_direction_probability,
            conditioning_values=val_conditioning,
            cfg=cfg,
        )

        test_prediction = _predict_quantiles(model, test_dl, cfg)
        test_prediction *= int(
            direction_calibration.get("direction_sign_multiplier", 1)
        )
        test_prediction = apply_weekly_sign_correction_np(
            test_prediction,
            float(direction_calibration.get("weekly_sign_threshold", 0.0)),
            horizon=cfg.forecast.primary_horizon_days,
        )
        if selected_direction_model.get("enabled"):
            test_prediction = apply_weekly_direction_model(
                test_prediction,
                test_direction_probability,
                threshold=float(
                    selected_direction_model.get("decision_threshold", 0.50)
                ),
                horizon=cfg.forecast.primary_horizon_days,
            )
        test_metrics = _compute_test_metrics_from_quantiles(
            test_actual,
            test_prediction,
            cfg,
            weekly_interval_scale=float(
                interval_calibration["weekly_interval_scale"]
            ),
            weekly_interval_calibration=interval_calibration,
            weekly_interval_conditioning_values=test_conditioning,
        )
        rank = _validation_rank(val_metrics, loss)
        validation_passed, validation_reasons = evaluate_quality_gate_metrics(
            val_metrics
        )
        test_passed, test_reasons = evaluate_quality_gate_metrics(test_metrics)
        rows.append(
            {
                "checkpoint": checkpoint.name,
                "checkpoint_loss": loss,
                "weekly_median_cap_factor": args.weekly_median_cap_factor,
                "effective_weekly_median_cap": effective_weekly_cap,
                "validation_rank": rank,
                "validation_gate_passed": validation_passed,
                "validation_gate_reasons": validation_reasons,
                "validation": val_metrics,
                "test_diagnostic_gate_passed": test_passed,
                "test_diagnostic_gate_reasons": test_reasons,
                "test_diagnostic": test_metrics,
            }
        )

    rows.sort(key=lambda row: row["validation_rank"])
    rendered = json.dumps(rows, indent=2, default=float)
    if args.output is not None:
        args.output.write_text(rendered + "\n", encoding="utf-8")
    else:
        print(rendered)


if __name__ == "__main__":
    main()