from miscore import Model, Process, Universe
from miscore.processes import Stratum
from miscore.utils import WeibullDistribution


class Clumsy(Process):
    # Process to divide smurfs into clumsy or regular
    def __init__(self, prob_clumsy: float = 0.5, name: str = "clumsy"):
        self.prob_clumsy = prob_clumsy
        self.name = name

    def properties(self, rng, n, properties):
        clumsy_category = rng.choice(
            ["clumsy", "regular"],
            size=n,
            p=[self.prob_clumsy, 1 - self.prob_clumsy]
        )
        return {
            "clumsy_category": clumsy_category
        }


class Smurf(Process):
    event_tags = [
        "eaten_early",
        "smurf_death"
    ]

    duration_tags = [
        "smurf_life"
    ]

    # Process to simulate smurf lifetimes
    def __init__(self, survival_shape: float = 2., survival_scale: float = 15,
                 name: str = "smurfs"):
        self.survival = WeibullDistribution(scale=survival_scale, shape=survival_shape)
        self.name = name

        self.callbacks = {
            "__start__": self.schedule_survival,
            "eaten_early": self.eaten_early,
            "__end__": self.log_durations
        }

    def properties(self, rng, n, properties):
        survival_r = rng.random(n)
        eaten_early_r = rng.random(n)
        return {
            "survival_r": survival_r,
            "eaten_early_r": eaten_early_r
        }

    def schedule_survival(self, individual):
        individual.add_to_queue_age(5., "eaten_early")
        natural_mortality_age = self.survival(individual.properties("survival_r"))
        individual.add_to_queue_age(natural_mortality_age, message="smurf_death",
                                    terminate=True)

    def eaten_early(self, individual):
        threshold = 0.1
        if individual.properties("clumsy_category") == "clumsy":
            threshold = 0.5
        if individual.properties("eaten_early_r") < threshold:
            individual.log_event("eaten_early")
            individual.add_to_queue(
                delta=0.,
                message="smurf_death",
                terminate=True
            )

    def log_durations(self, individual, message):
        individual.log_duration(
            tag="smurf_life",
            start_age=0,
            end_age=individual.age
        )


clumsy = Clumsy()

smurf = Smurf()

stratum = Stratum(
    [("clumsy_category", "==", "clumsy")]
)

universe = Universe(name="smurfs", processes=[clumsy, smurf, stratum])

model = Model(
    universes=[universe]
)

result = model.run(n=1000, seed=0)

print(result.events)
print(result.durations)


class Smurf_Screen(Process):
    event_tags = ["smurf_screen"]

    def __init__(self, test_sens: float = 0.5, cure_prob: float = 0.6):
        self.test_sens = test_sens
        self.cure_prob = cure_prob

        self.callbacks = {
            "__start__": self.schedule_screen,
            "smurf_screen": self.smurf_screen
        }

    def properties(self, rng, n, properties):
        return {
            "test_r": rng.random(n),
            "cure_r": rng.random(n)
        }

    def schedule_screen(self, individual):
        individual.add_to_queue_age(4, "smurf_screen")

    def smurf_screen(self, individual):
        individual.log_event("smurf_screen")
        if individual.properties("test_r") < self.test_sens:
            individual.hop_stratum(1)
            if individual.properties("clumsy_category"):
                if individual.properties("cure_r") < self.cure_prob:
                    individual.remove_from_queue(0)


smurf_screen = Smurf_Screen()

screening_universe = Universe(name="screened_smurfs", processes=[clumsy, smurf, smurf_screen])

model = Model(universes=[screening_universe])

result = model.run(n=1000, duration_years=[4], n_stratum_hops=1, seed=0)

print(result.events)
print(result.durations)

no_screening_universe = Universe(name="no_screening", processes=[clumsy, smurf])

model = Model([screening_universe, no_screening_universe])

result = model.run(n=1000, duration_years=[4], n_stratum_hops=1,
                   preserve_stratum="dynamic", seed=0)

print(result.events)
print(result.durations)

result = model.run(n=1000, duration_years=[4], n_stratum_hops=1,
                   preserve_stratum="final", seed=0)

print(result.events)
print(result.durations)