import argparse
import logging
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch
from torch_geometric.data import Batch

logging.getLogger("skfem").setLevel(logging.ERROR)

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

from src.network import create_model
from src.ppo import PPOTrainer
from src.utils import load_checkpoint, load_config, parse_center, save_checkpoint, setup_helmholtz_config
from src.visualize import visualize


def train(config: dict, iterations: int, checkpoint_dir: str = "checkpoints", save_freq: int = 50):
    t0 = time.time()
    algo = config.get("algorithm", {})
    dev = torch.device("cuda" if torch.cuda.is_available() and algo.get("use_gpu") else "cpu")
    print(f"[Device] {dev}")

    from environment.mesh_refinement import MeshRefinement

    env = MeshRefinement(
        environment_config=config.get("environment", {}).get("mesh_refinement", {}),
        seed=42,
    )
    print(f"[Env] node_feats={env.num_node_features} edge_feats={env.num_edge_features} act_dim={env.action_dimension}")

    model = create_model(env, config.get("network", {}), algo.get("ppo", {}), device=dev)
    print(f"[Model] params={sum(p.numel() for p in model.parameters()):,}")

    trainer = PPOTrainer(model, env, algo, device=dev)
    os.makedirs(checkpoint_dir, exist_ok=True)

    for it in range(1, iterations + 1):
        t1 = time.time()
        metrics = trainer.fit_iteration()
        print(
            f"  {it:4d}/{iterations} | loss={metrics['loss']:.4f} ev={metrics['explained_variance']:.3f} "
            f"agents={metrics['num_agents']:.0f} avg_r={metrics['avg_reward']:.4f} sum_r={metrics['sum_reward']:.2f} "
            f"x<0={metrics.get('neg_action_ratio', 0):.2f} "
            f"elig={metrics.get('eligible_ratio', 0):.2f} "
            f"sel={metrics.get('selected_count', 0):.0f} "
            f"{time.time() - t1:.1f}s"
        )
        if it % save_freq == 0 or it == iterations:
            save_checkpoint(model, model.optimizer, it, os.path.join(checkpoint_dir, f"model_iter{it:04d}.pt"))

    save_checkpoint(model, model.optimizer, iterations, os.path.join(checkpoint_dir, "model_final.pt"))
    print(f"[Train] done, total time {time.time() - t0:.1f}s")


def _eval_mie_error_test(env) -> float:
    """Compute relative L2 error of FEM vs Mie analytical solution."""
    fp = getattr(env.fem_problem, "fem_problem", None)
    if fp is None:
        return float("nan")
    _eps_r = getattr(fp, "_eps_r", None)
    _radius = getattr(fp, "_radius", None)
    _cx = getattr(fp, "_cx", None)
    _cy = getattr(fp, "_cy", None)
    _k = getattr(fp, "_k", None)
    if any(v is None for v in [_eps_r, _radius, _cx, _cy, _k]):
        return float("nan")

    from environment.mie_solution import mie_scattered_field
    pts = env.mesh.p.T
    u_mie = mie_scattered_field(pts, k0=_k, eps_r=_eps_r, radius=_radius, cx=_cx, cy=_cy)
    u_fem = env.scalar_solution
    diff = np.abs(u_fem - u_mie)
    denom = np.linalg.norm(np.abs(u_mie))
    if denom < 1e-12:
        denom = 1.0
    return float(np.linalg.norm(diff) / denom)


def test(config: dict, checkpoint_path: str, k_test=None, center=None, radius=None, eps_test=None):
    setup_helmholtz_config(config, k_test=k_test, center=center, radius=radius, eps_test=eps_test)
    algo = config.get("algorithm", {})

    from environment.mesh_refinement import MeshRefinement

    env = MeshRefinement(
        environment_config=config.get("environment", {}).get("mesh_refinement", {}),
        seed=99,
    )
    model = create_model(env, config.get("network", {}), algo.get("ppo", {}))
    load_checkpoint(model, checkpoint_path)
    model.eval()

    obs = env.reset()
    done = False
    step = 0
    n_elem_init = getattr(env, "_num_elements", env.num_agents)
    mie_err_0 = _eval_mie_error_test(env)
    print(f"  Step {step:2d}: reward=---      mie_err={mie_err_0:.4f}  elements={n_elem_init}"
          f"  budget={getattr(env, '_n_budget', '?')}")

    total_reward = 0.0
    while not done:
        with torch.no_grad():
            actions, _, _ = model(Batch.from_data_list([obs]), deterministic=True)
        obs, reward, done, info = env.step(actions.cpu().numpy())
        step_r = float(np.sum(reward))
        total_reward += step_r
        step += 1
        mie_err = _eval_mie_error_test(env)
        print(f"  Step {step:2d}: reward={step_r:+.4f} mie_err={mie_err:.4f}"
              f"  elements={info.get('num_elements', '?')} "
              f"x<0={info.get('neg_action_ratio', 0):.2f} sel={info.get('selected_count', 0)}")

    print(f"\n[Test] total_reward={total_reward:.4f} final_mie_error={mie_err:.4f}")


def main():
    parser = argparse.ArgumentParser(description="AFEM — Adaptive FEM with PPO RL")
    parser.add_argument("--mode", required=True, choices=["train", "test", "viz"])
    parser.add_argument("--config", default="src/config.yaml")
    parser.add_argument("--iterations", type=int, default=None)
    parser.add_argument("--checkpoint", default="checkpoints/model_final.pt")
    parser.add_argument("--checkpoint-dir", default="checkpoints")
    parser.add_argument("--save-freq", type=int, default=50)
    parser.add_argument("--output", default="result/visualization.png")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--k-test", type=float, default=None)
    parser.add_argument("--center", type=str, default=None)
    parser.add_argument("--radius", type=float, default=None)
    parser.add_argument("--eps-test", type=float, default=None)

    args = parser.parse_args()
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    cfg_path = args.config if os.path.isabs(args.config) else os.path.join(_project_root, args.config)
    config = load_config(cfg_path)
    if args.iterations is not None:
        config["iterations"] = args.iterations

    center = parse_center(args.center)

    if args.mode == "train":
        train(config, config.get("iterations", 100), args.checkpoint_dir, args.save_freq)
    elif args.mode == "test":
        test(config, args.checkpoint, k_test=args.k_test, center=center, radius=args.radius, eps_test=args.eps_test)
    elif args.mode == "viz":
        visualize(config, args.checkpoint, output_path=args.output, k_test=args.k_test, center=center, radius=args.radius, eps_test=args.eps_test)


if __name__ == "__main__":
    main()
