import unittest

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Subset

from dataset import AllFutureHealthDataset
from losses import ExponentialLoss, WeibullLoss
from models import DeepHealth
from train_util import fit_all_future_baseline


def _inverse_softplus(value: torch.Tensor) -> torch.Tensor:
    return torch.log(torch.expm1(value))


class AllFutureLikelihoodTests(unittest.TestCase):
    def test_exponential_uses_event_specific_first_onset_exposure(self):
        desired_rate = torch.tensor(
            [[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
            dtype=torch.float64,
        )
        logits = _inverse_softplus(desired_rate)
        criterion = ExponentialLoss(ignored_idx={0, 1, 2}, eps=1e-12)

        loss = criterion(
            logits=logits,
            targets=torch.tensor([[4, 0]]),
            exposure=torch.tensor([5.0], dtype=torch.float64),
            dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
            history=torch.tensor([[3, 0]]),
        )

        rate = F.softplus(logits) + criterion.eps
        expected = -(rate[0, 4].log()) + rate[0, 4] * 2.0 + rate[0, 5] * 5.0
        torch.testing.assert_close(loss, expected)

    def test_weibull_uses_event_time_and_excludes_prevalent_outcome(self):
        desired_rate = torch.tensor(
            [[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
            dtype=torch.float64,
        )
        logits = _inverse_softplus(desired_rate)
        rho = torch.tensor(
            [[1.0, 1.0, 1.0, 1.2, 1.5, 0.8]],
            dtype=torch.float64,
        )
        criterion = WeibullLoss(ignored_idx={0, 1, 2}, eps=1e-12)

        loss = criterion(
            logits=logits,
            weibull_rho=rho,
            targets=torch.tensor([[4, 0]]),
            dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
            exposure=torch.tensor([5.0], dtype=torch.float64),
            history=torch.tensor([[3, 0]]),
        )

        rate = F.softplus(logits) + criterion.eps
        log_hazard = (
            rate[0, 4].log()
            + rho[0, 4].log()
            + (rho[0, 4] - 1.0) * torch.tensor(2.0).log()
        )
        expected = (
            -log_hazard
            + rate[0, 4] * torch.pow(torch.tensor(2.0), rho[0, 4])
            + rate[0, 5] * torch.pow(torch.tensor(5.0), rho[0, 5])
        )
        torch.testing.assert_close(loss, expected)


class AllFutureQueryDistributionTests(unittest.TestCase):
    @staticmethod
    def _dataset() -> AllFutureHealthDataset:
        dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
        dataset.min_history_events = 1
        dataset.min_future_events = 1
        return dataset

    @staticmethod
    def _patient():
        return {
            "times": np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
            "labels": np.asarray([3, 4, 5, 6], dtype=np.int64),
            "t_obs": 4.0,
        }

    def test_train_validation_and_test_share_one_query_sampler(self):
        dataset = self._dataset()
        patient = self._patient()
        patient["query_intervals"] = dataset._eligible_query_intervals(patient)
        self.assertEqual(len(patient["query_intervals"]), 3)

        rng = np.random.RandomState(123)
        fixed = dataset._sample_fixed_validation_queries(patient, rng)
        self.assertEqual(len(fixed), 1)
        self.assertTrue(dataset._is_valid_query(patient, fixed[0]))

        seen_intervals = set()
        rng = np.random.RandomState(456)
        for _ in range(200):
            query = dataset.sample_query(patient, rng)
            self.assertTrue(dataset._is_valid_query(patient, query))
            seen_intervals.add(int(np.floor(query)))
        self.assertEqual(seen_intervals, {1, 2, 3})


class AllFutureBaselineTests(unittest.TestCase):
    def test_training_baseline_matches_first_onset_exposure(self):
        dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
        dataset.vocab_size = 7
        dataset.min_history_events = 1
        dataset.min_future_events = 2
        patient = {
            "times": np.asarray([1.0, 3.0, 5.0], dtype=np.float32),
            "labels": np.asarray([3, 4, 5], dtype=np.int64),
            "t_obs": 5.0,
            "query_intervals": [(1.0, 3.0)],
        }
        dataset.patients = [patient]
        subset = Subset(dataset, [0])

        stats = fit_all_future_baseline(dataset, subset, seed=17)
        query = dataset.sample_query(patient, np.random.RandomState(17))

        self.assertEqual(stats.event_count[3], 0)
        self.assertEqual(stats.at_risk_exposure[3], 0.0)
        self.assertEqual(stats.event_count[4], 1)
        self.assertEqual(stats.event_count[5], 1)
        self.assertAlmostEqual(stats.at_risk_exposure[4], 3.0 - query, places=5)
        self.assertAlmostEqual(stats.at_risk_exposure[5], 5.0 - query, places=5)
        self.assertAlmostEqual(
            float(F.softplus(torch.tensor(stats.bias[4]))),
            float(stats.rate[4]),
            places=6,
        )

    def test_model_starts_exactly_at_fitted_output_baseline(self):
        baseline_rate = torch.tensor([0.0, 0.0, 0.0, 0.02, 0.05, 0.10])
        baseline_bias = torch.zeros_like(baseline_rate)
        baseline_bias[3:] = _inverse_softplus(baseline_rate[3:])
        model = DeepHealth(
            vocab_size=6,
            n_embd=4,
            n_head=1,
            n_layer=1,
            n_types=1,
            n_cont_types=0,
            n_categories=1,
            cont_type_ids=[],
            target_mode="all_future",
            time_mode="absolute",
            dist_mode="weibull",
            model_architecture="transformer_ffn_v1",
            risk_head_bias=True,
            risk_head_bias_init=baseline_bias,
        )

        torch.testing.assert_close(model.risk_head.weight, torch.zeros_like(model.risk_head.weight))
        output_rate = F.softplus(model.risk_head(torch.randn(3, 4)))
        torch.testing.assert_close(output_rate[:, 3:], baseline_rate[None, 3:].expand(3, -1))
        torch.testing.assert_close(
            F.softplus(model.rho_head.bias),
            torch.ones_like(model.rho_head.bias),
            atol=2e-5,
            rtol=0.0,
        )


if __name__ == "__main__":
    unittest.main()