import torch
from mfkan import MFKAN, MFKANConfig, TrainingConfig, MFKANTrainer
from mfkan.utils import get_test_data

# Generate test data (or use your own)
data = get_test_data("test1", n_lf=50, n_hf=5)

# Configure model
config = MFKANConfig(
    input_dim=1,
    output_dim=1,
    lf_hidden_dims=[32, 32],
    nl_hidden_dims=[16, 16],
)

# Create and train
model = MFKAN(config)
trainer = MFKANTrainer(model, TrainingConfig(verbose=True))

# Phase 1: Train on abundant LF data
trainer.train_low_fidelity(data.x_lf, data.y_lf)

# Phase 2: Fine-tune on sparse HF data (LF frozen)
trainer.train_high_fidelity(data.x_hf, data.y_hf)

# Predict high-fidelity output
y_pred = model(data.x_test)
print(f"Learned α (nonlinear blend): {model.alpha.item():.3f}")

from mfkan import MFKANConfig

config = MFKANConfig(
    # Architecture
    input_dim=2,                    # Input dimension
    output_dim=1,                   # Output dimension
    lf_hidden_dims=[64, 64],        # K_L hidden layers
    nl_hidden_dims=[32, 32],        # K_nl hidden layers

    # KAN parameters
    grid_size=5,                    # B-spline grid intervals
    spline_order=3,                 # B-spline degree (k)
    grid_range=(-1.0, 1.0),         # Input normalization range

    # Alpha blending
    alpha_init=0.5,                 # Initial α value
    alpha_trainable=True,           # Learn α during HF training

    # Loss parameters
    lambda_alpha=0.01,              # α penalty weight
    alpha_exponent=4,               # α penalty power (n)
    regularization_weight=0.0,      # Spline regularization

    # Backend: "faster" (default), "efficient", or "pure"
    backend="faster",
)

from mfkan import TrainingConfig

config = TrainingConfig(
    # Epochs
    lf_epochs=500,                  # Low-fidelity training
    hf_epochs=500,                  # High-fidelity training

    # Learning rates
    lf_lr=1e-3,                     # LF learning rate
    hf_lr=1e-3,                     # HF learning rate

    # Batch sizes
    lf_batch_size=32,               # LF batch size
    hf_batch_size=16,               # HF batch size (often = N_HF)

    # Early stopping
    early_stopping_patience=0,       # 0 = disabled

    # Logging
    verbose=True,
    log_interval=100,
)

# Use faster-kan (recommended)
config = MFKANConfig(..., backend="faster")

# Fallback to pure PyTorch
config = MFKANConfig(..., backend="pure")

model = MFKAN(config)

# Forward passes
y_lf = model.forward_lf(x)          # Low-fidelity only
y_hf = model.forward_hf(x)          # High-fidelity (full model)
y = model(x)                        # Auto-select based on training state

# Properties
model.alpha                         # Current α value [0, 1]
model.is_lf_trained                 # LF phase completed?
model.is_lf_frozen                  # LF parameters frozen?

# Methods
model.freeze_low_fidelity()         # Freeze K_L after LF training
model.get_trainable_hf_params()     # Parameters for HF phase
model.regularization_loss()         # Spline regularization from K_nl

trainer = MFKANTrainer(model, config)

# Two-phase training
trainer.train_low_fidelity(x_lf, y_lf, val_data=None)
trainer.train_high_fidelity(x_hf, y_hf, val_data=None)

# Convenience method (both phases)
trainer.fit(x_lf, y_lf, x_hf, y_hf)

# Checkpoints
trainer.save_checkpoint("model.pt")
trainer.load_checkpoint("model.pt")

# Access training history
trainer.history["lf"]               # LF training metrics
trainer.history["hf"]               # HF training metrics

from mfkan.utils import compute_all_metrics, relative_l2_error

# All metrics at once
metrics = compute_all_metrics(y_pred, y_true)
# Returns: {relative_l2, mse, rmse, mae, max_error, r2}

# Individual metrics
error = relative_l2_error(y_pred, y_true)

from mfkan import MFKAN, MFKANConfig, TrainingConfig, MFKANTrainer
from mfkan.utils import generate_test1_jump_function, compute_all_metrics

# Data: f_L(x) = x, f_H(x) = x±1 (jump at x=0)
data = generate_test1_jump_function(n_lf=50, n_hf=5)

# Train MFKAN
model = MFKAN(MFKANConfig(input_dim=1, output_dim=1))
trainer = MFKANTrainer(model, TrainingConfig())
trainer.fit(data.x_lf, data.y_lf, data.x_hf, data.y_hf)

# Evaluate
y_pred = model(data.x_test)
metrics = compute_all_metrics(y_pred, data.y_test)
print(f"Relative L2 Error: {metrics['relative_l2']:.4f}")  # ~0.06
print(f"Final α: {model.alpha.item():.4f}")  # ~0 (linear)

from mfkan.utils import generate_test2_nonlinear_correlation

# Data: f_L(x) = sin(2πx), f_H(x) = sin²(2πx)
data = generate_test2_nonlinear_correlation(n_lf=51, n_hf=14)

# Nonlinear correlation requires larger K_nl
config = MFKANConfig(
    input_dim=1, output_dim=1,
    nl_hidden_dims=[32, 32],  # Larger for nonlinear
)
model = MFKAN(config)
# ... train as above

print(f"Final α: {model.alpha.item():.4f}")  # ~1 (nonlinear)