#!/usr/bin/env python3
"""Finite-time BP/FA diagnostic with measured time-varying tangent kernels."""

from __future__ import annotations

import argparse
import csv
import json
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import downstream_capacity_sweep as dcs  # noqa: E402
from fa_tangent_kernel_capacity import pseudo_jacobian  # noqa: E402


@dataclass(frozen=True)
class DiagnosticConfig:
    input_dim: int
    output_dim: int
    width: int
    train_samples: int
    test_samples: int
    steps: int
    lr: float
    init_seeds: int
    feedback_seeds: int
    data_seed: int
    feedback_scale: str
    device: str
    torch_threads: int
    outdir: str


@dataclass(frozen=True)
class DiagnosticRow:
    init_seed: int
    feedback_seed: int
    steps: int
    empirical_bp_loss: float
    empirical_fa_loss: float
    empirical_gap: float
    fixed_bp_loss: float
    fixed_fa_loss: float
    fixed_gap: float
    tv_bp_loss: float
    tv_fa_loss: float
    tv_gap: float
    fixed_gap_error: float
    tv_gap_error: float
    fixed_bp_error: float
    fixed_fa_error: float
    tv_bp_error: float
    tv_fa_error: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Measure finite-time fixedK vs time-varyingK predictions."
    )
    parser.add_argument("--input-dim", type=int, default=16)
    parser.add_argument("--output-dim", type=int, default=4)
    parser.add_argument("--width", type=int, default=64)
    parser.add_argument("--train-samples", type=int, default=128)
    parser.add_argument("--test-samples", type=int, default=512)
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--init-seeds", type=int, default=1)
    parser.add_argument("--feedback-seeds", type=int, default=2)
    parser.add_argument("--data-seed", type=int, default=123)
    parser.add_argument(
        "--feedback-scale",
        choices=["relu", "fan-in", "unit"],
        default="relu",
    )
    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    parser.add_argument("--torch-threads", type=int, default=8)
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/finite_time_kernel_diagnostic"),
    )
    return parser.parse_args()


def make_run_config(args: argparse.Namespace) -> dcs.RunConfig:
    return dcs.RunConfig(
        task="random",
        input_dim=args.input_dim,
        output_dim=args.output_dim,
        teacher_rank=4,
        teacher_width=64,
        teacher_hidden_layers=2,
        normalize_targets=False,
        widths=[args.width],
        train_samples=args.train_samples,
        test_samples=args.test_samples,
        probe_samples=8,
        steps=args.steps,
        lr=args.lr,
        optimizer="sgd",
        init_seeds=args.init_seeds,
        feedback_seeds=args.feedback_seeds,
        init_seed_offset=0,
        feedback_seed_offset=0,
        data_seed=args.data_seed,
        noise_std=0.0,
        feedback_scale=args.feedback_scale,
        capacity_q=0.01,
        jacobian_lambda_rel=1e-3,
        skip_jacobian=True,
        device=args.device,
        torch_threads=args.torch_threads,
        outdir=str(args.outdir),
        plot=False,
    )


def loss_from_residual(residual: np.ndarray, samples: int) -> float:
    return 0.5 * float(residual @ residual) / samples


def sgd_step(
    weights: list[torch.Tensor],
    x_train: torch.Tensor,
    y_train: torch.Tensor,
    lr: float,
    feedback: list[torch.Tensor] | None,
) -> list[torch.Tensor]:
    grads = dcs.gradients(weights, x_train, y_train, feedback)
    return [weight - lr * grad for weight, grad in zip(weights, grads)]


def kernel_for(
    weights: list[torch.Tensor],
    x_train: torch.Tensor,
    feedback: list[torch.Tensor] | None,
) -> np.ndarray:
    j_bp = pseudo_jacobian(weights, x_train, feedback=None).cpu().numpy()
    if feedback is None:
        return j_bp @ j_bp.T
    j_fa = pseudo_jacobian(weights, x_train, feedback=feedback).cpu().numpy()
    return j_bp @ j_fa.T


def run_one(
    config: dcs.RunConfig,
    x_train: torch.Tensor,
    y_train: torch.Tensor,
    init_seed: int,
    feedback_seed: int,
) -> DiagnosticRow:
    initial_weights = dcs.initialize_weights(config, config.widths[0], init_seed)
    feedback = dcs.init_feedback(config, config.widths[0], feedback_seed)
    with torch.no_grad():
        r0 = (dcs.predict(initial_weights, x_train) - y_train).reshape(-1).cpu().numpy()

    k_bp0 = kernel_for(initial_weights, x_train, feedback=None)
    k_fa0 = kernel_for(initial_weights, x_train, feedback=feedback)

    fixed_bp_residual = r0.copy()
    fixed_fa_residual = r0.copy()
    tv_bp_residual = r0.copy()
    tv_fa_residual = r0.copy()
    bp_weights = dcs.clone_weights(initial_weights)
    fa_weights = dcs.clone_weights(initial_weights)
    scale = config.lr / config.train_samples

    for _step in range(config.steps):
        fixed_bp_residual = fixed_bp_residual - scale * (k_bp0 @ fixed_bp_residual)
        fixed_fa_residual = fixed_fa_residual - scale * (k_fa0 @ fixed_fa_residual)

        k_bp_t = kernel_for(bp_weights, x_train, feedback=None)
        k_fa_t = kernel_for(fa_weights, x_train, feedback=feedback)
        tv_bp_residual = tv_bp_residual - scale * (k_bp_t @ tv_bp_residual)
        tv_fa_residual = tv_fa_residual - scale * (k_fa_t @ tv_fa_residual)

        bp_weights = sgd_step(bp_weights, x_train, y_train, config.lr, feedback=None)
        fa_weights = sgd_step(fa_weights, x_train, y_train, config.lr, feedback=feedback)

    empirical_bp_loss = dcs.mse(bp_weights, x_train, y_train)
    empirical_fa_loss = dcs.mse(fa_weights, x_train, y_train)
    fixed_bp_loss = loss_from_residual(fixed_bp_residual, config.train_samples)
    fixed_fa_loss = loss_from_residual(fixed_fa_residual, config.train_samples)
    tv_bp_loss = loss_from_residual(tv_bp_residual, config.train_samples)
    tv_fa_loss = loss_from_residual(tv_fa_residual, config.train_samples)

    empirical_gap = empirical_fa_loss - empirical_bp_loss
    fixed_gap = fixed_fa_loss - fixed_bp_loss
    tv_gap = tv_fa_loss - tv_bp_loss

    return DiagnosticRow(
        init_seed=init_seed,
        feedback_seed=feedback_seed,
        steps=config.steps,
        empirical_bp_loss=empirical_bp_loss,
        empirical_fa_loss=empirical_fa_loss,
        empirical_gap=empirical_gap,
        fixed_bp_loss=fixed_bp_loss,
        fixed_fa_loss=fixed_fa_loss,
        fixed_gap=fixed_gap,
        tv_bp_loss=tv_bp_loss,
        tv_fa_loss=tv_fa_loss,
        tv_gap=tv_gap,
        fixed_gap_error=fixed_gap - empirical_gap,
        tv_gap_error=tv_gap - empirical_gap,
        fixed_bp_error=fixed_bp_loss - empirical_bp_loss,
        fixed_fa_error=fixed_fa_loss - empirical_fa_loss,
        tv_bp_error=tv_bp_loss - empirical_bp_loss,
        tv_fa_error=tv_fa_loss - empirical_fa_loss,
    )


def write_rows(path: Path, rows: list[DiagnosticRow]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(DiagnosticRow.__annotations__.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))


def main() -> None:
    args = parse_args()
    if args.torch_threads > 0:
        torch.set_num_threads(args.torch_threads)
    config = make_run_config(args)
    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)

    rows: list[DiagnosticRow] = []
    total = args.init_seeds * args.feedback_seeds
    count = 0
    for init_index in range(args.init_seeds):
        init_seed = 10_000 + init_index
        for feedback_index in range(args.feedback_seeds):
            feedback_seed = 100_000 + init_index * 1000 + feedback_index
            count += 1
            print(
                f"[{count}/{total}] init={init_seed} feedback={feedback_seed}",
                flush=True,
            )
            rows.append(run_one(config, x_train, y_train, init_seed, feedback_seed))

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    write_rows(outdir / "finite_time_kernel_rows.csv", rows)
    payload = {
        "config": asdict(
            DiagnosticConfig(
                input_dim=args.input_dim,
                output_dim=args.output_dim,
                width=args.width,
                train_samples=args.train_samples,
                test_samples=args.test_samples,
                steps=args.steps,
                lr=args.lr,
                init_seeds=args.init_seeds,
                feedback_seeds=args.feedback_seeds,
                data_seed=args.data_seed,
                feedback_scale=args.feedback_scale,
                device=args.device,
                torch_threads=args.torch_threads,
                outdir=str(outdir),
            )
        ),
        "rows": [asdict(row) for row in rows],
    }
    (outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n")

    fixed_errors = np.array([row.fixed_gap_error for row in rows])
    tv_errors = np.array([row.tv_gap_error for row in rows])
    print(f"rows: {outdir / 'finite_time_kernel_rows.csv'}")
    print(
        "fixed gap error mean="
        f"{fixed_errors.mean():.6g}, mae={np.mean(np.abs(fixed_errors)):.6g}"
    )
    print(
        "time-varying gap error mean="
        f"{tv_errors.mean():.6g}, mae={np.mean(np.abs(tv_errors)):.6g}"
    )


if __name__ == "__main__":
    main()