#!/usr/bin/env python3
"""
Convert FNO weights from neuraloperator (PyTorch) to SafeTensors format.

This script can:
1. Train FNO models using neuraloperator and convert to SafeTensors
2. Convert existing PyTorch weights to SafeTensors
3. Generate random weights for testing

Usage:
    # Train and convert (Darcy Flow - built-in dataset)
    python convert_fno_weights.py --model fno2d_darcy --train --epochs 50

    # Convert existing weights
    python convert_fno_weights.py --model fno2d_darcy --input model.pt --output ./weights/

    # Generate random weights for testing
    python convert_fno_weights.py --model fno2d_darcy --random --output ./weights/

Requirements:
    Base: pip install torch safetensors numpy
    Training: pip install neuraloperator
"""

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Dict, Any, Optional

import numpy as np

try:
    import torch
    from safetensors.torch import save_file
except ImportError:
    print("Error: Required packages not installed.")
    print("Please run: pip install torch safetensors numpy")
    sys.exit(1)

# Conditional import for neuraloperator (only needed for training)
NEURALOP_AVAILABLE = False
try:
    from neuralop.models import FNO
    from neuralop.training import Trainer
    from neuralop.losses import H1Loss, LpLoss
    from neuralop.data.datasets import load_darcy_flow_small
    NEURALOP_AVAILABLE = True
except ImportError:
    pass


# Known FNO model configurations from neuraloperator
FNO_CONFIGS = {
    "fno2d_darcy": {
        "description": "FNO2d trained on Darcy Flow equation",
        "url": "https://github.com/neuraloperator/neuraloperator",
        "in_channels": 3,  # Lifting layer input: 1 data + 2 positional encoding
        "data_channels": 1,  # Raw data channels (used for neuraloperator model creation)
        "out_channels": 1,
        "width": 32,
        "n_modes": (12, 12),
        "n_layers": 4,
        "pde_type": "darcy",
    },
    "fno2d_navier_stokes": {
        "description": "FNO2d trained on 2D Navier-Stokes equation",
        "url": "https://github.com/neuraloperator/neuraloperator",
        "in_channels": 10,  # 10 time steps as input
        "out_channels": 1,
        "width": 64,
        "n_modes": (20, 20),
        "n_layers": 4,
        "pde_type": "navier_stokes",
    },
    "fno2d_poisson": {
        "description": "FNO2d trained on Poisson equation",
        "url": "https://github.com/neuraloperator/neuraloperator",
        "in_channels": 1,
        "out_channels": 1,
        "width": 32,
        "n_modes": (12, 12),
        "n_layers": 4,
        "pde_type": "poisson",
    },
}


def create_random_fno_weights(config: Dict[str, Any], seed: int = 42, v2_format: bool = True) -> Dict[str, torch.Tensor]:
    """
    Create random FNO weights matching the neuraloperator architecture.

    This is used for testing when pre-trained weights are not available.

    Args:
        config: Model configuration
        seed: Random seed for reproducibility
        v2_format: If True, use neuraloperator v2.0 format with lifting MLP
    """
    torch.manual_seed(seed)
    np.random.seed(seed)

    # For v2.0 format, use data_channels for the actual input
    # data_channels = raw input channels, in_channels = data_channels + 2 (positional encoding)
    data_channels = config.get("data_channels", config["in_channels"])
    out_channels = config["out_channels"]
    width = config["width"]
    n_modes = config["n_modes"]
    n_layers = config["n_layers"]

    weights = {}

    if v2_format:
        # neuraloperator v2.0 format with lifting MLP
        # Lifting MLP: (data_channels + 2) -> 2*width -> width
        # The +2 is for positional encoding (x, y coordinates)
        lifting_in = data_channels + 2
        hidden_dim = 2 * width

        # First layer: (data_channels + 2) -> 2*width
        weights["lifting.fcs.0.weight"] = torch.randn(hidden_dim, lifting_in) * 0.02
        weights["lifting.fcs.0.bias"] = torch.zeros(hidden_dim)

        # Second layer: 2*width -> width
        weights["lifting.fcs.1.weight"] = torch.randn(width, hidden_dim) * 0.02
        weights["lifting.fcs.1.bias"] = torch.zeros(width)
    else:
        # Legacy format: single lifting layer
        in_channels = config["in_channels"]
        weights["fc0.weight"] = torch.randn(width, in_channels) * 0.02
        weights["fc0.bias"] = torch.zeros(width)

    # Spectral convolution layers and skip convolutions
    for i in range(n_layers):
        # SpectralConv2d weights (complex)
        # Shape: [in_channels, out_channels, modes1, modes2, 2]
        # The 2 is for real and imaginary parts
        spectral_shape = (width, width, n_modes[0], n_modes[1])

        # Real and imaginary parts stored separately in neuraloperator
        weights[f"conv{i}.weights1"] = torch.randn(*spectral_shape, 2) * 0.02
        weights[f"conv{i}.weights2"] = torch.randn(*spectral_shape, 2) * 0.02

        # Skip connection (1x1 conv implemented as Linear in some versions)
        weights[f"w{i}.weight"] = torch.randn(width, width) * 0.02
        weights[f"w{i}.bias"] = torch.zeros(width)

    # Projection layers (fc1, fc2): width -> 128 -> out_channels
    weights["fc1.weight"] = torch.randn(128, width) * 0.02
    weights["fc1.bias"] = torch.zeros(128)

    weights["fc2.weight"] = torch.randn(out_channels, 128) * 0.02
    weights["fc2.bias"] = torch.zeros(out_channels)

    return weights


def convert_neuraloperator_weights(pt_weights: Dict[str, torch.Tensor],
                                    config: Dict[str, Any]) -> Dict[str, torch.Tensor]:
    """
    Convert neuraloperator weight naming to rtx-neural-operator naming convention.

    Supports two formats:
    1. Manual/test format (fc0, conv{i}, w{i}, fc1, fc2 or lifting.fcs.*)
    2. neuraloperator library format (lifting.fcs, fno_blocks.convs, projection.fcs)

    rtx-neural-operator v2.0 naming:
        lifting.fcs.0.weight, lifting.fcs.0.bias (first layer of lifting MLP)
        lifting.fcs.1.weight, lifting.fcs.1.bias (second layer of lifting MLP)
        spectral_conv.{i}.weights1_real, spectral_conv.{i}.weights1_imag
        spectral_conv.{i}.weights2_real, spectral_conv.{i}.weights2_imag
        conv.{i}.weight, conv.{i}.bias
        projection.0.weight, projection.0.bias
        projection.1.weight, projection.1.bias
    """
    converted = {}
    n_layers = config.get("n_layers", 4)

    # Detect format based on key presence
    is_library_format = any(k.startswith("fno_blocks.") for k in pt_weights.keys())
    is_v2_lifting = any(k.startswith("lifting.fcs.") for k in pt_weights.keys())

    if is_library_format:
        # neuraloperator library format (v2.0+)
        # Weights are Conv1d format: [out_channels, in_channels, kernel_size=1]
        # We squeeze the kernel dimension for Linear format: [out_channels, in_channels]

        # Lifting MLP: keep as lifting.fcs.0/1.weight/bias for v2.0 format
        for layer_idx in [0, 1]:
            w_key = f"lifting.fcs.{layer_idx}.weight"
            b_key = f"lifting.fcs.{layer_idx}.bias"
            if w_key in pt_weights:
                w = pt_weights[w_key]
                # Conv1d weight [out, in, 1] -> Linear weight [out, in]
                converted[w_key] = w.squeeze(-1) if w.dim() == 3 else w
            if b_key in pt_weights:
                converted[b_key] = pt_weights[b_key]

        # Spectral convolutions: fno_blocks.convs.{i}.weight.tensor (complex, single tensor)
        for i in range(n_layers):
            # neuraloperator v2.0 uses .weight.tensor for complex spectral weights
            w_key = f"fno_blocks.convs.{i}.weight.tensor"

            if w_key in pt_weights:
                w = pt_weights[w_key]
                # Shape: [in_ch, out_ch, modes_h, modes_w] as complex64
                # Split into real and imaginary parts
                if torch.is_complex(w):
                    # Note: neuraloperator uses a single weight tensor for both halves
                    # We map it to weights1 for rtx-neural-operator
                    converted[f"spectral_conv.{i}.weights1_real"] = w.real.contiguous()
                    converted[f"spectral_conv.{i}.weights1_imag"] = w.imag.contiguous()
                    # Create matching weights2 (same as weights1 for compatibility)
                    converted[f"spectral_conv.{i}.weights2_real"] = w.real.contiguous()
                    converted[f"spectral_conv.{i}.weights2_imag"] = w.imag.contiguous()
                else:
                    converted[f"spectral_conv.{i}.weights1_real"] = w
                    converted[f"spectral_conv.{i}.weights1_imag"] = torch.zeros_like(w)
                    converted[f"spectral_conv.{i}.weights2_real"] = w.clone()
                    converted[f"spectral_conv.{i}.weights2_imag"] = torch.zeros_like(w)

            # Skip connections: fno_blocks.fno_skips.{i}.conv.weight (no bias in v2.0)
            skip_w_key = f"fno_blocks.fno_skips.{i}.conv.weight"

            if skip_w_key in pt_weights:
                w = pt_weights[skip_w_key]
                # Conv1d weight [out, in, 1] -> Linear weight [out, in]
                converted[f"conv.{i}.weight"] = w.squeeze(-1) if w.dim() == 3 else w
                # Create zero bias for compatibility with rtx-neural-operator
                converted[f"conv.{i}.bias"] = torch.zeros(w.shape[0])

        # Projection layers: projection.fcs.0/1.weight/bias
        if "projection.fcs.0.weight" in pt_weights:
            w = pt_weights["projection.fcs.0.weight"]
            converted["projection.0.weight"] = w.squeeze(-1) if w.dim() == 3 else w
        if "projection.fcs.0.bias" in pt_weights:
            converted["projection.0.bias"] = pt_weights["projection.fcs.0.bias"]
        if "projection.fcs.1.weight" in pt_weights:
            w = pt_weights["projection.fcs.1.weight"]
            converted["projection.1.weight"] = w.squeeze(-1) if w.dim() == 3 else w
        if "projection.fcs.1.bias" in pt_weights:
            converted["projection.1.bias"] = pt_weights["projection.fcs.1.bias"]

    else:
        # Manual/test format (fc0, conv{i}, w{i}, fc1, fc2) or v2 random format
        # Lifting layer - check for v2 format first
        if is_v2_lifting:
            # v2 format: lifting.fcs.0/1.weight/bias (2-layer MLP)
            for layer_idx in [0, 1]:
                w_key = f"lifting.fcs.{layer_idx}.weight"
                b_key = f"lifting.fcs.{layer_idx}.bias"
                if w_key in pt_weights:
                    converted[w_key] = pt_weights[w_key]
                if b_key in pt_weights:
                    converted[b_key] = pt_weights[b_key]
        elif "fc0.weight" in pt_weights:
            # Legacy format: single linear layer
            converted["lifting.weight"] = pt_weights["fc0.weight"]
            if "fc0.bias" in pt_weights:
                converted["lifting.bias"] = pt_weights["fc0.bias"]

        # Spectral convolutions and skip connections
        for i in range(n_layers):
            # Spectral conv weights (complex -> real + imag)
            w1_key = f"conv{i}.weights1"
            w2_key = f"conv{i}.weights2"

            if w1_key in pt_weights:
                w1 = pt_weights[w1_key]
                if w1.shape[-1] == 2:  # Has real/imag dimension
                    converted[f"spectral_conv.{i}.weights1_real"] = w1[..., 0]
                    converted[f"spectral_conv.{i}.weights1_imag"] = w1[..., 1]
                else:
                    converted[f"spectral_conv.{i}.weights1_real"] = w1.real if torch.is_complex(w1) else w1
                    converted[f"spectral_conv.{i}.weights1_imag"] = w1.imag if torch.is_complex(w1) else torch.zeros_like(w1)

            if w2_key in pt_weights:
                w2 = pt_weights[w2_key]
                if w2.shape[-1] == 2:
                    converted[f"spectral_conv.{i}.weights2_real"] = w2[..., 0]
                    converted[f"spectral_conv.{i}.weights2_imag"] = w2[..., 1]
                else:
                    converted[f"spectral_conv.{i}.weights2_real"] = w2.real if torch.is_complex(w2) else w2
                    converted[f"spectral_conv.{i}.weights2_imag"] = w2.imag if torch.is_complex(w2) else torch.zeros_like(w2)

            # Skip connection (1x1 conv)
            w_key = f"w{i}.weight"
            b_key = f"w{i}.bias"

            if w_key in pt_weights:
                converted[f"conv.{i}.weight"] = pt_weights[w_key]
            if b_key in pt_weights:
                converted[f"conv.{i}.bias"] = pt_weights[b_key]

        # Projection layers
        if "fc1.weight" in pt_weights:
            converted["projection.0.weight"] = pt_weights["fc1.weight"]
        if "fc1.bias" in pt_weights:
            converted["projection.0.bias"] = pt_weights["fc1.bias"]
        if "fc2.weight" in pt_weights:
            converted["projection.1.weight"] = pt_weights["fc2.weight"]
        if "fc2.bias" in pt_weights:
            converted["projection.1.bias"] = pt_weights["fc2.bias"]

    return converted


def save_weights_safetensors(weights: Dict[str, torch.Tensor],
                              output_path: Path,
                              metadata: Optional[Dict[str, str]] = None) -> None:
    """Save weights in SafeTensors format."""
    # Ensure all tensors are contiguous and on CPU
    clean_weights = {}
    for name, tensor in weights.items():
        if tensor.is_cuda:
            tensor = tensor.cpu()
        if not tensor.is_contiguous():
            tensor = tensor.contiguous()
        # Convert to float32 for compatibility
        if tensor.dtype != torch.float32:
            tensor = tensor.float()
        clean_weights[name] = tensor

    save_file(clean_weights, output_path, metadata=metadata)
    print(f"Saved weights to {output_path}")


def save_config(config: Dict[str, Any], output_path: Path) -> None:
    """Save model configuration as JSON."""
    with open(output_path, "w") as f:
        json.dump(config, f, indent=2)
    print(f"Saved config to {output_path}")


def train_fno_model(
    config: Dict[str, Any],
    n_epochs: int,
    batch_size: int,
    learning_rate: float,
    n_train: int,
    n_test: int,
    device: str,
) -> Dict[str, torch.Tensor]:
    """
    Train an FNO model using neuraloperator's Trainer API.

    Currently supports Darcy Flow (built-in dataset).

    Returns the trained model's state_dict.
    """
    if not NEURALOP_AVAILABLE:
        raise RuntimeError(
            "neuraloperator is required for training.\n"
            "Install with: pip install neuraloperator"
        )

    pde_type = config["pde_type"]
    if pde_type != "darcy":
        raise ValueError(
            f"Training for '{pde_type}' is not yet supported.\n"
            f"Only 'darcy' (Darcy Flow) has a built-in dataset."
        )

    # 1. Determine device
    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

    device_obj = torch.device(device)
    print(f"Training on device: {device_obj}")

    # 2. Load Darcy Flow dataset
    print(f"Loading Darcy Flow dataset (n_train={n_train}, n_test={n_test})...")
    train_loader, test_loaders, data_processor = load_darcy_flow_small(
        n_train=n_train,
        batch_size=batch_size,
        n_tests=[n_test],
        test_resolutions=[32],
        test_batch_sizes=[batch_size],
    )

    # 3. Create model
    # neuraloperator uses hidden_channels instead of width
    # Use data_channels (raw data) for model creation - model adds positional encoding internally
    data_channels = config.get("data_channels", config["in_channels"])
    print(f"Creating FNO model (data_channels={data_channels})...")
    model = FNO(
        n_modes=config["n_modes"],
        hidden_channels=config["width"],
        in_channels=data_channels,
        out_channels=config["out_channels"],
        n_layers=config["n_layers"],
    )
    model = model.to(device_obj)

    # Print model info
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {total_params:,}")

    # 4. Setup optimizer and scheduler
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=learning_rate,
        weight_decay=1e-4
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=n_epochs
    )

    # 5. Setup loss functions
    train_loss = H1Loss(d=2)  # H1 loss includes gradient information
    eval_losses = {
        'h1': H1Loss(d=2),
        'l2': LpLoss(d=2, p=2),
    }

    # 6. Create trainer
    trainer = Trainer(
        model=model,
        n_epochs=n_epochs,
        data_processor=data_processor,
        device=device_obj,
        verbose=True,
    )

    # 7. Print training configuration
    print("\n" + "=" * 60)
    print("Training Configuration")
    print("=" * 60)
    print(f"  PDE Type: {config['pde_type']}")
    print(f"  In Channels: {config['in_channels']}")
    print(f"  Out Channels: {config['out_channels']}")
    print(f"  Width: {config['width']}")
    print(f"  Modes: {config['n_modes']}")
    print(f"  Layers: {config['n_layers']}")
    print("-" * 60)
    print(f"  Epochs: {n_epochs}")
    print(f"  Batch Size: {batch_size}")
    print(f"  Learning Rate: {learning_rate}")
    print(f"  Training Samples: {n_train}")
    print(f"  Test Samples: {n_test}")
    print("=" * 60 + "\n")

    # 8. Train
    print(f"Starting training for {n_epochs} epochs...")
    trainer.train(
        train_loader=train_loader,
        test_loaders=test_loaders,
        optimizer=optimizer,
        scheduler=scheduler,
        regularizer=False,
        training_loss=train_loss,
        eval_losses=eval_losses,
    )

    # 9. Final evaluation
    print("\nFinal evaluation:")
    # test_loaders is a dict keyed by resolution
    test_resolution = list(test_loaders.keys())[0]
    final_metrics = trainer.evaluate(eval_losses, test_loaders[test_resolution])
    for name, value in final_metrics.items():
        print(f"  {name}: {value:.6f}")

    # 10. Return state dict
    return model.state_dict()


def list_models() -> None:
    """List available FNO model configurations."""
    print("\nAvailable FNO models:")
    print("-" * 60)
    for name, config in FNO_CONFIGS.items():
        print(f"\n{name}:")
        print(f"  Description: {config['description']}")
        print(f"  PDE Type: {config['pde_type']}")
        print(f"  In Channels: {config['in_channels']}")
        print(f"  Out Channels: {config['out_channels']}")
        print(f"  Width: {config['width']}")
        print(f"  Modes: {config['n_modes']}")
        print(f"  Layers: {config['n_layers']}")
        if pde_type := config.get('pde_type'):
            trainable = "Yes (built-in dataset)" if pde_type == "darcy" else "Requires external data"
            print(f"  Trainable: {trainable}")


def main():
    parser = argparse.ArgumentParser(
        description="Convert FNO weights from PyTorch to SafeTensors format"
    )
    parser.add_argument(
        "--model",
        type=str,
        choices=list(FNO_CONFIGS.keys()) + ["all"],
        help="Model to convert (use 'all' for all models)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="./weights/fno",
        help="Output directory for converted weights"
    )
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="Path to PyTorch weights file (.pt or .pth)"
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available models"
    )
    parser.add_argument(
        "--random",
        action="store_true",
        help="Generate random weights (for testing)"
    )
    # Training arguments
    parser.add_argument(
        "--train",
        action="store_true",
        help="Train the FNO model before converting (requires neuraloperator)"
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=50,
        help="Number of training epochs (default: 50)"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=16,
        help="Training batch size (default: 16)"
    )
    parser.add_argument(
        "--learning-rate",
        type=float,
        default=1e-3,
        help="Learning rate (default: 1e-3)"
    )
    parser.add_argument(
        "--n-train",
        type=int,
        default=1000,
        help="Number of training samples (default: 1000)"
    )
    parser.add_argument(
        "--n-test",
        type=int,
        default=100,
        help="Number of test samples (default: 100)"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        choices=["auto", "cuda", "cpu", "mps"],
        help="Device for training (default: auto-detect)"
    )

    args = parser.parse_args()

    if args.list:
        list_models()
        return

    if not args.model:
        parser.print_help()
        print("\nError: --model is required")
        return

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    models_to_convert = list(FNO_CONFIGS.keys()) if args.model == "all" else [args.model]

    for model_name in models_to_convert:
        config = FNO_CONFIGS[model_name]
        print(f"\nProcessing {model_name}...")

        model_dir = output_dir / model_name
        model_dir.mkdir(parents=True, exist_ok=True)

        if args.train:
            # Train the model using neuraloperator
            if not NEURALOP_AVAILABLE:
                print("Error: neuraloperator is required for training.")
                print("Install with: pip install neuraloperator")
                return
            print(f"Training {model_name}...")
            pt_weights = train_fno_model(
                config=config,
                n_epochs=args.epochs,
                batch_size=args.batch_size,
                learning_rate=args.learning_rate,
                n_train=args.n_train,
                n_test=args.n_test,
                device=args.device,
            )
        elif args.input and os.path.exists(args.input):
            # Load from provided PyTorch file
            print(f"Loading weights from {args.input}")
            pt_weights = torch.load(args.input, map_location="cpu", weights_only=True)
            if "state_dict" in pt_weights:
                pt_weights = pt_weights["state_dict"]
            elif "model" in pt_weights:
                pt_weights = pt_weights["model"]
        elif args.random:
            # Generate random weights for testing
            print("Generating random weights for testing...")
            pt_weights = create_random_fno_weights(config)
        else:
            print(f"Note: No pre-trained weights provided. Generating random weights.")
            print(f"To train a model, use --train flag (requires neuraloperator)")
            print(f"To convert existing weights, use --input")
            pt_weights = create_random_fno_weights(config)

        # Convert weight naming convention
        converted_weights = convert_neuraloperator_weights(pt_weights, config)

        # Print weight info
        print(f"Converted {len(converted_weights)} weight tensors:")
        total_params = 0
        for name, tensor in converted_weights.items():
            params = tensor.numel()
            total_params += params
            print(f"  {name}: {list(tensor.shape)} ({params:,} params)")
        print(f"Total parameters: {total_params:,}")

        # Save as SafeTensors
        weight_source = "trained" if args.train else ("converted" if args.input else "random")
        metadata = {
            "model_name": model_name,
            "pde_type": config["pde_type"],
            "framework": "rtx-neural-operator",
            "source": "neuraloperator",
            "weight_source": weight_source,
        }
        if args.train:
            metadata["epochs"] = str(args.epochs)
            metadata["n_train"] = str(args.n_train)

        weights_path = model_dir / "model.safetensors"
        save_weights_safetensors(converted_weights, weights_path, metadata)

        # Save config
        config_path = model_dir / "config.json"
        save_config(config, config_path)

        print(f"Model {model_name} saved to {model_dir}")


if __name__ == "__main__":
    main()