import dataclasses
import re
import tempfile

import numpy as np
import pytest
import torch
import xarray as xr

from fme.ace.inference.evaluator import main as inference_evaluator_main
from fme.ace.testing import DimSizes, MonthlyReferenceData, save_scalar_netcdf
from fme.ace.testing.fv3gfs_data import get_nd_dataset
from fme.ace.train.train import main as train_main
from fme.core.coordinates import LatLonCoordinates
from fme.core.testing.wandb import mock_wandb


def save_ocean_nd_netcdf(
    filename: str,
    dim_sizes: DimSizes,
    variable_names: list[str],
    timestep_days: float = 1.0,
    nz_levels: int = 2,  # Number of ocean levels for thetao_0, mask_0 etc.
):
    """
    Saves a netCDF file with synthetic ocean data, including masks and NaNs.
    Inspired by fme.ace.testing.fv3gfs_data.save_nd_netcdf.
    """
    base_variable_names = [
        name
        for name in variable_names
        if not name.startswith("mask_") and not name.startswith("idepth_")
    ]

    # Use a modified DimSizes for base dataset to avoid issues with nz_interface
    # if get_nd_dataset interprets it in a way not suitable for ocean vars
    base_dim_sizes = dataclasses.replace(dim_sizes, nz_interface=1)

    ds = get_nd_dataset(
        dim_sizes=base_dim_sizes,
        variable_names=base_variable_names,
        timestep_days=timestep_days,
        include_vertical_coordinate=False,  # No ak/bk for ocean
    )

    rng = np.random.default_rng(0)  # For reproducibility

    horizontal_dims = list(map(lambda x: x.name, dim_sizes.horizontal))
    horizontal_shape = list(map(lambda x: x.size, dim_sizes.horizontal))

    masks_data = {}
    for i in range(nz_levels):
        mask_name = f"mask_{i}"
        if mask_name in variable_names:
            mask_data_values = rng.integers(low=0, high=2, size=horizontal_shape)
            coords_for_mask = {
                dim: ds.coords[dim] for dim in horizontal_dims if dim in ds.coords
            }
            masks_data[i] = xr.DataArray(
                mask_data_values,
                dims=horizontal_dims,
                coords=coords_for_mask,
                name=mask_name,
            )
            ds[mask_name] = masks_data[i]

    if "mask_0" in ds:
        ds["mask_2d"] = ds["mask_0"].copy()

    for var_name in base_variable_names:
        if var_name not in ds:  # Skip if not generated by get_nd_dataset
            continue
        match = re.fullmatch(r"(.+)_(\d+)", var_name)
        if match:
            level = int(match.group(2))
            if level < nz_levels and level in masks_data:
                ds[var_name] = ds[var_name].where(masks_data[level] == 1, float("nan"))
        else:  # Assume 2D variable, apply mask_0 if it exists (e.g., for sst)
            if 0 in masks_data and "mask_0" in variable_names:
                ds[var_name] = ds[var_name].where(masks_data[0] == 1, float("nan"))

    # Add idepth variables (interface depths)
    # nz_interface in DimSizes should be nz_levels + 1
    if dim_sizes.nz_interface != nz_levels + 1:
        raise ValueError(
            f"dim_sizes.nz_interface ({dim_sizes.nz_interface}) "
            f"must be nz_levels ({nz_levels}) + 1."
        )
    for i in range(dim_sizes.nz_interface):
        idepth_name = f"idepth_{i}"
        if idepth_name in variable_names:
            ds[idepth_name] = float(i)  # Saved as 0-D variable

    ds.to_netcdf(filename, unlimited_dims=["time"])
    return ds


_TRAIN_CONFIG_TEMPLATE = """
experiment_dir: {results_dir}
save_checkpoint: true
save_per_epoch_diagnostics: {save_per_epoch_diagnostics}
max_epochs: {max_epochs}
n_forward_steps: 2
logging:
  log_to_screen: true
  log_to_wandb: {log_to_wandb}
  log_to_file: false
  project: fme
  entity: ai2cm
train_loader:
  dataset:
    - data_path: '{train_data_path}'
      spatial_dimensions: latlon
  batch_size: 2
  num_data_workers: 0
validation_loader:
  dataset:
    - data_path: '{valid_data_path}'
      spatial_dimensions: latlon
  batch_size: 2
  num_data_workers: 0
optimization:
  use_gradient_accumulation: true
  optimizer_type: "Adam"
  lr: 0.001
  kwargs:
    weight_decay: 0.01
  scheduler:
      type: CosineAnnealingLR
      kwargs:
        T_max: 1
stepper:
  loss:
    type: "MSE"
  input_masking:
    mask_value: 0
    fill_value: 0.0
  step:
    type: single_module
    config:
      in_names: {in_variable_names}
      out_names: {out_variable_names}
      normalization:
        network:
          global_means_path: '{global_means_path}'
          global_stds_path: '{global_stds_path}'
      builder:
        type: Samudra
        config:
          ch_width: [8, 16]
          dilation: [2, 4]
          n_layers: [1, 1]
          norm: batch
          norm_kwargs:
            track_running_stats: false
      corrector:
        type: "ocean_corrector"
        config:
          sea_ice_fraction_correction:
            sea_ice_fraction_name: sea_ice_fraction
            land_fraction_name: land_fraction
          ocean_heat_content_correction: true
inference:
  aggregator:
    monthly_reference_data: {monthly_data_filename}
  loader:
    dataset:
      data_path: '{valid_data_path}'
      spatial_dimensions: latlon
    start_indices:
      first: 0
      n_initial_conditions: 2
      interval: 1
  n_forward_steps: {inference_forward_steps}
  forward_steps_in_memory: 2
"""

_INFERENCE_CONFIG_TEMPLATE = """
experiment_dir: {results_dir}
n_forward_steps: 6
forward_steps_in_memory: 2
checkpoint_path: {results_dir}/training_checkpoints/best_ckpt.tar
data_writer:
  save_prediction_files: true
aggregator:
  log_video: true
logging:
  log_to_screen: true
  log_to_wandb: {log_to_wandb}
  log_to_file: false
  project: fme
  entity: ai2cm
loader:
  dataset:
    data_path: '{valid_data_path}'
    spatial_dimensions: latlon
    fill_nans:
      method: constant
      value: 0.0
  start_indices:
    first: 0
    n_initial_conditions: 2
    interval: 1
"""


def _get_test_yaml_files(
    train_data_path,
    valid_data_path,
    monthly_data_filename,
    results_dir,
    global_means_path,
    global_stds_path,
    in_variable_names,
    out_variable_names,
    log_to_wandb=False,
    max_epochs=1,
    inference_forward_steps=2,
    save_per_epoch_diagnostics=False,
):
    train_config_content = _TRAIN_CONFIG_TEMPLATE

    train_string = train_config_content.format(
        train_data_path=train_data_path,
        valid_data_path=valid_data_path,
        monthly_data_filename=monthly_data_filename,
        results_dir=results_dir,
        global_means_path=global_means_path,
        global_stds_path=global_stds_path,
        in_variable_names=in_variable_names,
        out_variable_names=out_variable_names,
        log_to_wandb=str(log_to_wandb).lower(),
        max_epochs=max_epochs,
        inference_forward_steps=inference_forward_steps,
        save_per_epoch_diagnostics=str(save_per_epoch_diagnostics).lower(),
    )

    inference_string = _INFERENCE_CONFIG_TEMPLATE.format(
        results_dir=results_dir,
        log_to_wandb=str(log_to_wandb).lower(),
        valid_data_path=valid_data_path,  # Inference uses validation data path
    )

    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".yaml") as f_train:
        f_train.write(train_string)

    with tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=".yaml"
    ) as f_inference:
        f_inference.write(inference_string)

    return f_train.name, f_inference.name


def get_ocean_dim_sizes(
    n_lat=16,
    n_lon=32,
    n_time=3,
    nz_levels=2,  # For thetao_0, thetao_1, mask_0, mask_1
) -> DimSizes:
    return DimSizes(
        n_time=n_time,
        horizontal=LatLonCoordinates(
            lon=torch.arange(float(n_lon)),
            lat=torch.arange(float(n_lat)),
        ).loaded_sizes,
        nz_interface=nz_levels + 1,  # For idepth_0, idepth_1, idepth_2
    )


def _setup(
    path,
    log_to_wandb=True,
    max_epochs=1,
    n_time=60,
    timestep_days=2,
    inference_forward_steps=10,
    save_per_epoch_diagnostics=True,
    nz_levels=2,
):
    if not path.exists():
        path.mkdir(parents=True, exist_ok=True)
    seed = 0
    np.random.seed(seed)
    torch.manual_seed(seed)

    # Ocean-specific variable names
    in_variable_names = [
        "thetao_0",
        "thetao_1",
        "sst",
        "hfds",
        "hfgeou",
        "sea_surface_fraction",
        "sea_ice_fraction",
        "land_fraction",
    ]
    out_variable_names = [
        "thetao_0",
        "thetao_1",
        "sst",
        "sea_ice_fraction",
    ]

    # Add masks and idepths for data generation
    mask_names = [f"mask_{i}" for i in range(nz_levels)]
    idepth_names = [f"idepth_{i}" for i in range(nz_levels + 1)]

    in_variable_names_for_data_gen = list(set(in_variable_names + mask_names))

    all_variable_names_for_data_gen = list(
        set(in_variable_names_for_data_gen + out_variable_names + idepth_names)
    )
    stats_variable_names = list(set(in_variable_names + out_variable_names))

    dim_sizes = get_ocean_dim_sizes(n_time=n_time, nz_levels=nz_levels)

    data_dir = path / "data"
    stats_dir = path / "stats"
    results_dir = path / "results"
    data_dir.mkdir(exist_ok=True)
    stats_dir.mkdir(exist_ok=True)
    results_dir.mkdir(exist_ok=True)

    save_ocean_nd_netcdf(
        filename=data_dir / "data.nc",
        dim_sizes=dim_sizes,
        variable_names=all_variable_names_for_data_gen,
        timestep_days=timestep_days,
        nz_levels=nz_levels,
    )
    save_scalar_netcdf(
        stats_dir / "stats-mean.nc",
        variable_names=stats_variable_names,
    )
    save_scalar_netcdf(
        stats_dir / "stats-stddev.nc",
        variable_names=stats_variable_names,
    )

    monthly_dim_sizes = get_ocean_dim_sizes(n_time=10 * 12, nz_levels=1)
    monthly_out_vars = [
        v for v in out_variable_names if not re.match(r".+_\d+", v) or v.endswith("_0")
    ]
    if not monthly_out_vars and out_variable_names:
        monthly_out_vars = [out_variable_names[0]]

    if monthly_out_vars:
        monthly_reference_data = MonthlyReferenceData(
            path=data_dir,
            names=monthly_out_vars,
            dim_sizes=monthly_dim_sizes,
            n_ensemble=3,
        )
        monthly_data_filename = monthly_reference_data.data_filename
    else:
        monthly_data_filename = "null"

    train_config_filename, inference_config_filename = _get_test_yaml_files(
        train_data_path=data_dir,
        valid_data_path=data_dir,
        monthly_data_filename=monthly_data_filename,
        results_dir=results_dir,
        global_means_path=stats_dir / "stats-mean.nc",
        global_stds_path=stats_dir / "stats-stddev.nc",
        in_variable_names=in_variable_names,
        out_variable_names=out_variable_names,
        log_to_wandb=log_to_wandb,
        max_epochs=max_epochs,
        inference_forward_steps=inference_forward_steps,
        save_per_epoch_diagnostics=save_per_epoch_diagnostics,
    )
    return train_config_filename, inference_config_filename


def test_train_and_inference(tmp_path, very_fast_only: bool):
    """Ensure that ACE Ocean training and subsequent standalone inference run."""
    if very_fast_only:
        pytest.skip("Skipping non-fast tests")

    train_config, inference_config = _setup(tmp_path)

    with mock_wandb() as wandb:
        train_main(yaml_config=train_config)
        wandb_logs = wandb.get_logs()
        for log in wandb_logs:
            assert "inference/mean/forecast_step" not in log

    results_dir = tmp_path / "results"

    validation_output_dir = results_dir / "output" / "val" / "epoch_0001"
    assert validation_output_dir.exists()
    validation_diags = ["mean"]
    validation_map_diags = ["snapshot", "mean_map"]
    for diagnostic in validation_diags + validation_map_diags:
        diagnostic_output = validation_output_dir / f"{diagnostic}_diagnostics.nc"
        assert diagnostic_output.exists()
        ds = xr.open_dataset(diagnostic_output, decode_timedelta=False)
        assert len(ds) > 0
        for var in ds.data_vars:
            if diagnostic == "mean":
                assert not np.isnan(ds[var].values).any()
            elif "ocean_heat_content" in var:
                assert not np.isnan(ds[var].values).all()
                assert np.isnan(ds[var].values).any()
            else:
                # other
                assert not np.isnan(ds[var].values).all()
                assert np.isnan(ds[var].values).any()

    inline_inference_output_dir = results_dir / "output" / "inference" / "epoch_0001"
    assert inline_inference_output_dir.exists()
    for diagnostic in ("time_mean", "time_mean_norm"):
        diagnostic_output = inline_inference_output_dir / f"{diagnostic}_diagnostics.nc"
        assert diagnostic_output.exists()
        ds = xr.open_dataset(diagnostic_output, decode_timedelta=False)
        assert len(ds) > 0
        for var in ds.data_vars:
            assert not np.isnan(ds[var].values).all()
            assert np.isnan(ds[var].values).any()

    with mock_wandb() as wandb:
        wandb.configure(log_to_wandb=True)
        inference_evaluator_main(yaml_config=inference_config)
        wandb_logs = wandb.get_logs()

    tm_logs = wandb_logs[-1]
    for name, metric in tm_logs.items():
        if "power_spectrum" in name:
            # power spectra not well defined for masked data
            continue
        if isinstance(metric, float):
            assert not np.isnan(metric), f"{name} should not be NaN"

    prediction_output_path = results_dir / "autoregressive_predictions.nc"
    best_checkpoint_path = results_dir / "training_checkpoints" / "best_ckpt.tar"
    assert prediction_output_path.exists()
    assert best_checkpoint_path.exists()

    ds_prediction = xr.open_dataset(prediction_output_path, decode_timedelta=False)
    for name in ["sst", "thetao_0", "thetao_1", "ocean_heat_content"]:
        assert name in ds_prediction
        # outputs should have some non-null values
        assert not np.isnan(ds_prediction[name].values).all()
        # outputs should have some masked regions
        assert np.isnan(ds_prediction[name].values).any()

    ds_target = xr.open_dataset(
        results_dir / "autoregressive_target.nc", decode_timedelta=False
    )
    assert "sst" in ds_target

    for diagnostic in ("mean", "mean_norm"):
        diagnostic_output = results_dir / f"{diagnostic}_diagnostics.nc"
        assert diagnostic_output.exists()
        ds = xr.open_dataset(diagnostic_output, decode_timedelta=False)
        assert len(ds) > 0
        for var in ds.data_vars:
            assert not np.isnan(ds[var].values).any()
