from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

from example_common import (
    diagnostics,
    forecast_ensemble,
    global_latent_localized_ensrf,
    inflate_ensemble,
    load_checkpoint,
    make_initial_ensemble,
    make_params,
    make_truth,
)
from lda.ks import integrate_steps
from lda.ks_ae import KSGlobalAutoencoder
from lda.utils import ensure_dir


def main() -> None:
    parser = argparse.ArgumentParser(description="Example global-latent EnSRF with decoded physical localization.")
    parser.add_argument("--n", type=int, default=512)
    parser.add_argument("--length", type=float, default=None)
    parser.add_argument("--dt", type=float, default=0.05)
    parser.add_argument("--cycle-steps", type=int, default=40)
    parser.add_argument("--cycles", type=int, default=10)
    parser.add_argument("--members", type=int, default=20)
    parser.add_argument("--obs-stride", type=int, default=32)
    parser.add_argument("--obs-error", type=float, default=0.5)
    parser.add_argument("--inflation", type=float, default=1.10)
    parser.add_argument("--loc-radius", type=float, default=192.0)
    parser.add_argument("--latent-dim", type=int, default=256)
    parser.add_argument("--bottleneck-grid", type=int, default=64)
    parser.add_argument("--base-channels", type=int, default=16)
    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--spinup-steps", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=3)
    parser.add_argument("--output", type=Path, default=Path("example_outputs/gla_ensrf_loc.npz"))
    args = parser.parse_args()

    params = make_params(args.n, args.length)
    rng = np.random.default_rng(args.seed)
    truth = make_truth(params, rng, args.spinup_steps, args.dt)
    ensemble = make_initial_ensemble(truth, params, rng, args.members, perturbation_std=0.5)
    obs_locs = np.arange(0, args.n, args.obs_stride)

    model = KSGlobalAutoencoder(
        n=args.n,
        latent_dim=args.latent_dim,
        bottleneck_grid=args.bottleneck_grid,
        base_channels=args.base_channels,
    )
    model = load_checkpoint(model, args.checkpoint)
    norm_mean = float(ensemble.mean())
    norm_std = float(max(ensemble.std(), 1.0e-6))

    prior_rmse = []
    analysis_rmse = []
    for _ in range(args.cycles):
        truth = integrate_steps(truth, args.cycle_steps, args.dt, params)
        ensemble = forecast_ensemble(ensemble, args.cycle_steps, args.dt, params)
        prior = ensemble.copy()
        obs = truth[obs_locs] + rng.normal(0.0, args.obs_error, size=len(obs_locs))
        ensemble = inflate_ensemble(ensemble, args.inflation)
        ensemble = global_latent_localized_ensrf(
            ensemble,
            obs,
            obs_locs,
            args.obs_error,
            args.loc_radius,
            model,
            norm_mean,
            norm_std,
        )
        d = diagnostics(prior, ensemble, truth)
        prior_rmse.append(d[0])
        analysis_rmse.append(d[1])

    ensure_dir(args.output.parent)
    np.savez_compressed(
        args.output,
        prior_rmse=np.asarray(prior_rmse),
        analysis_rmse=np.asarray(analysis_rmse),
        obs_locs=obs_locs,
        inflation=args.inflation,
        loc_radius=args.loc_radius,
        latent_dim=args.latent_dim,
    )
    print(f"saved {args.output}")


if __name__ == "__main__":
    main()
