"""Custom ODE template with TODOs for user to fill in."""

from anypinn.cli._types import DataSource
from anypinn.cli.templates._base import train_py_core, train_py_lightning

EXPERIMENT_NAME = "custom-ode"


def _ode_py(data_source: DataSource) -> str:
    if data_source == DataSource.CSV:
        validation_block = """\
validation: ValidationRegistry = {
    # TODO: add validation sources
    # "param_name": ColumnRef(column="your_column"),
}"""
    else:
        validation_block = """\
validation: ValidationRegistry = {
    # TODO: add validation sources
    # "param_name": lambda x: torch.full_like(x, TRUE_VALUE),
}"""

    return f'''\
"""Custom ODE — mathematical definition."""

from __future__ import annotations

from typing import cast

import torch
from torch import Tensor

from anypinn.core import (
    ArgsRegistry,
    Argument,
    ColumnRef,
    Field,
    FieldsRegistry,
    Parameter,
    ParamsRegistry,
    ValidationRegistry,
)
from anypinn.lightning.callbacks import DataScaling
from anypinn.problems import ODEHyperparameters, ODEInverseProblem, ODEProperties

# ============================================================================
# Keys — define string keys for your state variables and parameters
# ============================================================================

# TODO: define state variable keys
# Y1_KEY = "y1"
# Y2_KEY = "y2"

# TODO: define parameter keys
# PARAM_KEY = "param"

# ============================================================================
# Constants
# ============================================================================

# TODO: set your initial conditions
# Y1_0 = 1.0
# Y2_0 = 0.0

# TODO: set time domain
# T_TOTAL = 10

# ============================================================================
# ODE Definition
# ============================================================================


def my_ode(x: Tensor, y: Tensor, args: ArgsRegistry) -> Tensor:
    """TODO: implement your ODE system.

    Args:
        x: Independent variable (e.g. time). Shape: (N,)
        y: State variables. Each element has shape (N,)
        args: Parameters (both fixed and learnable).

    Returns:
        dy/dx stacked as a tensor.
    """
    # Example:
    # y1, y2 = y
    # p = args[PARAM_KEY]
    # dy1 = -p(x) * y1
    # dy2 = p(x) * y1 - y2
    # return torch.stack([dy1, dy2])
    raise NotImplementedError("TODO: implement your ODE")


# ============================================================================
# Validation
# ============================================================================

{validation_block}

# ============================================================================
# Data Module Factory
# ============================================================================


def create_data_module(hp: ODEHyperparameters):
    # TODO: create and return your DataModule
    # See anypinn.catalog for examples of DataModule subclasses.
    raise NotImplementedError("TODO: implement create_data_module")


# ============================================================================
# Problem Factory
# ============================================================================


def create_problem(hp: ODEHyperparameters) -> ODEInverseProblem:
    # TODO: define your problem
    # props = ODEProperties(
    #     ode=my_ode,
    #     y0=torch.tensor([Y1_0, Y2_0]),
    #     args={{
    #         # fixed arguments go here
    #     }},
    # )
    #
    # fields = FieldsRegistry({{
    #     Y1_KEY: Field(config=hp.fields_config),
    #     Y2_KEY: Field(config=hp.fields_config),
    # }})
    # params = ParamsRegistry({{
    #     PARAM_KEY: Parameter(config=hp.params_config),
    # }})
    #
    # def predict_data(
    #     x_data: Tensor, fields: FieldsRegistry, _params: ParamsRegistry
    # ) -> Tensor:
    #     return cast(Tensor, fields[Y1_KEY](x_data))
    #
    # return ODEInverseProblem(
    #     props=props, hp=hp, fields=fields, params=params,
    #     predict_data=predict_data,
    # )
    raise NotImplementedError("TODO: implement create_problem")
'''


def _config_py(data_source: DataSource) -> str:
    if data_source == DataSource.CSV:
        training_data = """\
    training_data=IngestionConfig(
        batch_size=100,
        data_ratio=2,
        collocations=6000,
        df_path=Path("./data/data.csv"),
        y_columns=["y_obs"],  # TODO: update column names
    ),"""
    else:
        training_data = """\
    training_data=GenerationConfig(
        batch_size=100,
        data_ratio=2,
        collocations=6000,
        x=torch.linspace(start=0, end=10, steps=100),
        args_to_train={},
        noise_level=0,
    ),"""

    return f'''\
"""Custom ODE — training configuration."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch

from anypinn.core import (
    GenerationConfig,
    IngestionConfig,
    MLPConfig,
    ScalarConfig,
    SchedulerConfig,
)
from anypinn.problems import ODEHyperparameters


# ============================================================================
# Run Configuration
# ============================================================================


@dataclass
class RunConfig:
    max_epochs: int
    gradient_clip_val: float
    experiment_name: str
    run_name: str


CONFIG = RunConfig(
    max_epochs=2000,
    gradient_clip_val=0.1,
    experiment_name="{EXPERIMENT_NAME}",
    run_name="v0",
)

# ============================================================================
# Hyperparameters — TODO: tune these for your problem
# ============================================================================

hp = ODEHyperparameters(
    lr=1e-3,
{training_data}
    fields_config=MLPConfig(
        in_dim=1,
        out_dim=1,
        hidden_layers=[64, 128, 128, 64],
        activation="tanh",
        output_activation=None,
    ),
    params_config=ScalarConfig(
        init_value=0.5,
    ),
    scheduler=SchedulerConfig(
        mode="min",
        factor=0.5,
        patience=55,
        threshold=5e-3,
        min_lr=1e-6,
    ),
    pde_weight=1,
    ic_weight=1,
    data_weight=1,
)
'''


def render(data_source: DataSource, lightning: bool) -> dict[str, str]:
    train_fn = train_py_lightning if lightning else train_py_core
    return {
        "ode.py": _ode_py(data_source),
        "config.py": _config_py(data_source),
        "train.py": train_fn(EXPERIMENT_NAME),
    }
