import random
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import cast

import pytest
import torch
from pydantic import field_validator
from torch import Tensor, nn
from torch.optim import SGD

from refiners.fluxion import layers as fl
from refiners.fluxion.utils import norm
from refiners.training_utils.callback import Callback, CallbackConfig
from refiners.training_utils.clock import ClockConfig
from refiners.training_utils.common import (
    Epoch,
    Iteration,
    Step,
    TimeValue,
    TimeValueInput,
    count_learnable_parameters,
    human_readable_number,
    parse_number_unit_field,
    scoped_seed,
)
from refiners.training_utils.config import BaseConfig, ModelConfig
from refiners.training_utils.trainer import (
    Trainer,
    TrainingClock,
    WarmupScheduler,
    count_learnable_parameters,
    human_readable_number,
    register_callback,
    register_model,
)


@dataclass
class MockBatch:
    inputs: torch.Tensor
    targets: torch.Tensor


class MockModelConfig(ModelConfig):
    use_activation: bool


class MockCallbackConfig(CallbackConfig):
    on_batch_end_interval: Step | Iteration | Epoch
    on_batch_end_seed: int
    on_optimizer_step_interval: Iteration | Epoch

    @field_validator("on_batch_end_interval", "on_optimizer_step_interval", mode="before")
    def parse_field(cls, value: TimeValueInput) -> TimeValue:
        return parse_number_unit_field(value)


class MockConfig(BaseConfig):
    # we register the `early_callback` before the `clock` callback to test the callback ordering
    early_callback: CallbackConfig = CallbackConfig()
    clock: ClockConfig = ClockConfig()

    mock_model: MockModelConfig
    mock_callback: MockCallbackConfig


class MockModel(fl.Chain):
    def __init__(self):
        super().__init__(
            fl.Linear(10, 10),
            fl.Linear(10, 10),
            fl.Linear(10, 10),
        )

    def add_activation(self) -> None:
        self.insert(1, fl.SiLU())
        self.insert(3, fl.SiLU())


class MockCallback(Callback["MockTrainer"]):
    def __init__(self, config: MockCallbackConfig) -> None:
        self.config = config
        self.optimizer_step_count = 0
        self.step_end_count = 0
        self.optimizer_step_random_int: int | None = None
        self.step_end_random_int: int | None = None

    def on_init_begin(self, trainer: "MockTrainer") -> None:
        pass

    def on_optimizer_step_begin(self, trainer: "MockTrainer") -> None:
        if not trainer.clock.is_due(self.config.on_optimizer_step_interval):
            return
        self.optimizer_step_count += 1
        self.optimizer_step_random_int = random.randint(0, 100)

    def on_step_end(self, trainer: "MockTrainer") -> None:
        if not trainer.clock.is_due(self.config.on_batch_end_interval):
            return

        # We verify that this callback is always called before the clock is updated (see `_call_callbacks` in trainer.py)
        assert trainer.clock.step // 3 <= self.step_end_count

        self.step_end_count += 1
        with scoped_seed(self.config.on_batch_end_seed):
            self.step_end_random_int = random.randint(0, 100)


class EarlyMockCallback(Callback["MockTrainer"]):
    """
    A callback that will be registered before the Clock callback to test the callback ordering.
    """

    def on_train_begin(self, trainer: "MockTrainer") -> None:
        assert trainer.clock.start_time is not None, "Clock callback should have been called before this callback."


class MockTrainer(Trainer[MockConfig, MockBatch]):
    step_counter: int = 0
    model_registration_counter: int = 0

    @property
    def dataset_length(self) -> int:
        return 20

    def get_item(self, index: int) -> MockBatch:
        return MockBatch(inputs=torch.randn(1, 10), targets=torch.randn(1, 10))

    def collate_fn(self, batch: list[MockBatch]) -> MockBatch:
        return MockBatch(
            inputs=torch.cat([b.inputs for b in batch]),
            targets=torch.cat([b.targets for b in batch]),
        )

    @register_callback()
    def early_callback(self, config: CallbackConfig) -> EarlyMockCallback:
        return EarlyMockCallback()

    @register_callback()
    def mock_callback(self, config: MockCallbackConfig) -> MockCallback:
        return MockCallback(config)

    @register_model()
    def mock_model(self, config: MockModelConfig) -> MockModel:
        model = MockModel()
        if config.use_activation:
            model.add_activation()

        self.model_registration_counter += 1
        return model

    def compute_loss(self, batch: MockBatch) -> Tensor:
        self.step_counter += 1
        inputs, targets = batch.inputs.to(self.device), batch.targets.to(self.device)
        outputs = self.mock_model(inputs)
        return norm(outputs - targets)


@pytest.fixture
def mock_config() -> MockConfig:
    config = MockConfig.load_from_toml(Path(__file__).parent / "mock_config.toml")
    return config


@pytest.fixture
def mock_trainer(mock_config: MockConfig) -> MockTrainer:
    return MockTrainer(config=mock_config)


@pytest.fixture
def mock_trainer_short(mock_config: MockConfig) -> MockTrainer:
    mock_config_short = mock_config.model_copy(deep=True)
    mock_config_short.training.duration = Step(3)
    return MockTrainer(config=mock_config_short)


@pytest.fixture
def mock_model() -> fl.Chain:
    return MockModel()


def test_count_learnable_parameters_with_params() -> None:
    params = [
        nn.Parameter(torch.randn(2, 2), requires_grad=True),
        nn.Parameter(torch.randn(5), requires_grad=False),
        nn.Parameter(torch.randn(3, 3), requires_grad=True),
    ]
    # cast because of PyTorch 2.2, see https://github.com/pytorch/pytorch/issues/118736
    assert count_learnable_parameters(cast(list[nn.Parameter], params)) == 13


def test_count_learnable_parameters_with_model(mock_model: fl.Chain) -> None:
    assert count_learnable_parameters(mock_model.parameters()) == 330


def test_human_readable_number() -> None:
    assert human_readable_number(123) == "123.0"
    assert human_readable_number(1234) == "1.2K"
    assert human_readable_number(1234567) == "1.2M"


@pytest.fixture
def training_clock() -> TrainingClock:
    return TrainingClock(
        batch_size=10,
        training_duration=Epoch(5),
        gradient_accumulation=Step(1),
        lr_scheduler_interval=Epoch(1),
    )


def test_zero_batch_size_error():
    with pytest.raises(AssertionError):
        TrainingClock(
            batch_size=0,
            training_duration=Epoch(5),
            gradient_accumulation=Step(1),
            lr_scheduler_interval=Epoch(1),
        )


def test_timer_functionality(training_clock: TrainingClock) -> None:
    training_clock.start_timer()
    assert training_clock.start_time is not None
    training_clock.stop_timer()
    assert training_clock.end_time is not None
    assert training_clock.time_elapsed >= 0


def test_mock_trainer_initialization(mock_config: MockConfig, mock_trainer: MockTrainer) -> None:
    assert mock_trainer.config == mock_config
    assert isinstance(mock_trainer, MockTrainer)
    assert mock_trainer.optimizer is not None
    assert mock_trainer.lr_scheduler is not None
    assert mock_trainer.model_registration_counter == 1


def test_training_cycle(mock_trainer: MockTrainer) -> None:
    clock = mock_trainer.clock
    config = mock_trainer.config

    assert mock_trainer.step_counter == 0
    assert clock.epoch == 0

    mock_trainer.train()

    assert clock.epoch == config.training.duration.number
    assert mock_trainer.step_counter == mock_trainer.clock.step


def test_callback_registration(mock_trainer: MockTrainer) -> None:
    mock_trainer.train()

    # Check that the callback skips every other iteration
    assert mock_trainer.mock_callback.optimizer_step_count == mock_trainer.clock.iteration // 2
    assert mock_trainer.mock_callback.step_end_count == mock_trainer.clock.step // 3 + 1

    # Check that the random seed was set
    assert mock_trainer.mock_callback.optimizer_step_random_int == 93
    assert mock_trainer.mock_callback.step_end_random_int == 81


def test_training_short_cycle(mock_trainer_short: MockTrainer) -> None:
    clock = mock_trainer_short.clock
    config = mock_trainer_short.config

    assert mock_trainer_short.step_counter == 0
    assert mock_trainer_short.clock.epoch == 0

    mock_trainer_short.train()

    assert clock.step == config.training.duration.number


@pytest.fixture
def warmup_scheduler():
    optimizer = SGD([nn.Parameter(torch.randn(2, 2), requires_grad=True)], lr=0.1)
    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, 1)
    return WarmupScheduler(optimizer, scheduler, warmup_scheduler_steps=100)


def test_initial_lr(warmup_scheduler: WarmupScheduler) -> None:
    optimizer = warmup_scheduler.optimizer
    for group in optimizer.param_groups:
        assert group["lr"] == 1e-3


def test_warmup_lr(warmup_scheduler: WarmupScheduler) -> None:
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            category=UserWarning,
            message=r"Detected call of `lr_scheduler.step\(\)` before `optimizer.step\(\)`",
        )
        for _ in range(102):
            warmup_scheduler.step()
    optimizer = warmup_scheduler.optimizer
    for group in optimizer.param_groups:
        assert group["lr"] == 0.1


class MockTrainerWith2Models(MockTrainer):
    @register_model()
    def mock_model1(self, config: ModelConfig) -> MockModel:
        return MockModel()

    @register_model()
    def mock_model2(self, config: ModelConfig) -> MockModel:
        return MockModel()

    def compute_loss(self, batch: MockBatch) -> Tensor:
        self.step_counter += 1
        inputs, targets = batch.inputs.to(self.device), batch.targets.to(self.device)
        outputs = self.mock_model2(self.mock_model1(inputs))
        return norm(outputs - targets)


class MockConfig_2_Models(BaseConfig):
    mock_model1: ModelConfig
    mock_model2: ModelConfig


@pytest.fixture
def mock_config_2_models() -> MockConfig_2_Models:
    return MockConfig_2_Models.load_from_toml(Path(__file__).parent / "mock_config_2_models.toml")


@pytest.fixture
def mock_trainer_2_models(mock_config_2_models: MockConfig) -> MockTrainerWith2Models:
    return MockTrainerWith2Models(config=mock_config_2_models)


def test_optimizer_parameters(mock_trainer_2_models: MockTrainerWith2Models) -> None:
    assert len(mock_trainer_2_models.optimizer.param_groups) == 2
    assert mock_trainer_2_models.optimizer.param_groups[0]["lr"] == 1e-5