from pathlib import Path
from typing import Union

from eq.data import Catalog, InMemoryDataset, default_catalogs_dir
from eq.models import ETAS


class ETAS_MultiCatalog(Catalog):
    """Multiple catalogs generated by the ETAS model."""

    def __init__(
        self,
        root_dir: Union[str, Path] = default_catalogs_dir / "ETAS_MultiCatalog",
        num_sequences: int = 1000,
        train_frac: float = 0.6,
        val_frac: float = 0.2,
        test_frac: float = 0.2,
        t_end: float = 10_000,
        max_length: int = 20_000,
        base_rate: float = 0.02,
        omori_p: float = 1.08,
        omori_c: float = 0.1,
        productivity_k: float = 0.0073,
        productivity_alpha: float = 1.0,
        richter_b: float = 1.0,
        mag_completeness: float = 2.0,
        random_state: int = 123,
    ):
        metadata = {
            "name": "ETAS_MultiCatalog",
            "num_sequences": num_sequences,
            "train_frac": train_frac,
            "val_frac": val_frac,
            "test_frac": test_frac,
            "t_end": t_end,
            "max_length": max_length,
            "base_rate": base_rate,
            "omori_p": omori_p,
            "omori_c": omori_c,
            "productivity_k": productivity_k,
            "productivity_alpha": productivity_alpha,
            "richter_b": richter_b,
            "mag_completeness": mag_completeness,
            "random_state": random_state,
        }
        super().__init__(root_dir=root_dir, metadata=metadata)

        self.train = InMemoryDataset.load_from_disk(self.root_dir / "train.pt")
        self.val = InMemoryDataset.load_from_disk(self.root_dir / "val.pt")
        self.test = InMemoryDataset.load_from_disk(self.root_dir / "test.pt")

    @property
    def required_files(self):
        return ["train.pt", "val.pt", "test.pt", "metadata.pt"]

    def generate_catalog(self):
        etas_model = ETAS(
            omori_p_init=self.metadata["omori_p"],
            omori_c_init=self.metadata["omori_c"],
            base_rate_init=self.metadata["base_rate"],
            productivity_k_init=self.metadata["productivity_k"],
            productivity_alpha_init=self.metadata["productivity_alpha"],
            richter_b=self.metadata["richter_b"],
            mag_completeness=self.metadata["mag_completeness"],
        )
        sequences = etas_model.sample(
            batch_size=self.metadata["num_sequences"],
            duration=self.metadata["t_end"],
            random_state=self.metadata["random_state"],
            max_length=self.metadata["max_length"],
            return_sequences=True,
        )
        num_train = int(len(sequences) * self.metadata["train_frac"])
        num_val = int(len(sequences) * self.metadata["val_frac"])
        d_train = InMemoryDataset(sequences=sequences[:num_train])
        d_val = InMemoryDataset(sequences=sequences[num_train : num_train + num_val])
        d_test = InMemoryDataset(sequences=sequences[num_train + num_val :])
        d_train.save_to_disk(self.root_dir / "train.pt")
        d_val.save_to_disk(self.root_dir / "val.pt")
        d_test.save_to_disk(self.root_dir / "test.pt")


class ETAS_SingleCatalog(Catalog):
    def __init__(
        self,
        root_dir: Union[Path, str] = default_catalogs_dir / "ETAS_SingleCatalog",
        t_end: float = 3_000_000,
        max_length: int = 100_000,
        t_val_start: float = 2_000_000,
        t_test_start: float = 2_500_000,
        random_state: int = 1,
        base_rate: float = 0.02,
        omori_p: float = 1.08,
        omori_c: float = 0.1,
        productivity_k: float = 0.005,
        productivity_alpha: float = 0.8,
        richter_b: float = 1.0,
        mag_completeness: float = 2.0,
    ):
        metadata = {
            "name": "ETAS_SingleCatalog",
            "t_end": t_end,
            "max_length": max_length,
            "base_rate": base_rate,
            "omori_p": omori_p,
            "omori_c": omori_c,
            "productivity_k": productivity_k,
            "productivity_alpha": productivity_alpha,
            "richter_b": richter_b,
            "mag_completeness": mag_completeness,
            "random_state": random_state,
        }
        super().__init__(root_dir=root_dir, metadata=metadata)

        # Split the sequence into train, test and val parts
        self.full_sequence = InMemoryDataset.load_from_disk(
            self.root_dir / "full_sequence.pt"
        )[0]
        t_start = self.full_sequence.t_start
        seq_train = self.full_sequence.get_subsequence(t_start, t_val_start)
        seq_val = self.full_sequence.get_subsequence(t_start, t_test_start)
        seq_val.t_nll_start = t_val_start
        seq_test = self.full_sequence.get_subsequence(t_start, self.full_sequence.t_end)
        seq_test.t_nll_start = t_test_start

        self.train = InMemoryDataset([seq_train])
        self.val = InMemoryDataset([seq_val])
        self.test = InMemoryDataset([seq_test])

    @property
    def required_files(self):
        return ["full_sequence.pt", "metadata.pt"]

    def generate_catalog(self):
        etas_model = ETAS(
            omori_p_init=self.metadata["omori_p"],
            omori_c_init=self.metadata["omori_c"],
            base_rate_init=self.metadata["base_rate"],
            productivity_k_init=self.metadata["productivity_k"],
            productivity_alpha_init=self.metadata["productivity_alpha"],
            richter_b=self.metadata["richter_b"],
            mag_completeness=self.metadata["mag_completeness"],
        )
        sequences = etas_model.sample(
            batch_size=1,
            duration=self.metadata["t_end"],
            random_state=self.metadata["random_state"],
            max_length=None,
            return_sequences=True,
        )
        dataset = InMemoryDataset(sequences=sequences)
        dataset.save_to_disk(self.root_dir / "full_sequence.pt")


class ToyETAS:
    def __init__(self):
        dset = InMemoryDataset.load_from_disk(default_catalogs_dir / "etas_kelian.pt")
        self.train = dset
        self.val = dset
        self.test = dset
        self.metadata = {
            "base_rate": 0.02,
            "omori_p": 1.08,
            "omori_c": 0.1,
            "productivity_k": 0.0073,
            "productivity_alpha": 1.0,
            "richter_b": 1.0,
            "mag_completeness": 2.0,
        }
