from __future__ import annotations

from collections import Counter
import datetime as dt
import json
import math
from queue import PriorityQueue
from typing import Union, Optional, Tuple, cast, TYPE_CHECKING, Dict

import numpy as np

from slim import logger, LoggableMixin
from slim.simulation.config import Config
from slim.types.TreatmentTypes import GeneticMechanism, Money, ChemicalTreatment, \
    ThermalTreatment
from slim.simulation.lice_population import (GenoDistrib, GrossLiceDistrib,
                                             LicePopulation, GenoTreatmentDistrib, GenoTreatmentValue,
                                             GenoLifeStageDistrib, largest_remainder, LifeStage, GenoDistribDict)
from slim.types.QueueTypes import (DamAvailabilityBatch, EggBatch, TravellingEggBatch, TreatmentEvent,
                                   pop_from_queue)
from slim.JSONEncoders import CustomFarmEncoder

if TYPE_CHECKING:  # pragma: no cover
    from slim.simulation.farm import Farm, GenoDistribByHatchDate

OptionalDamBatch = Optional[DamAvailabilityBatch]
OptionalEggBatch = Optional[EggBatch]

class Cage(LoggableMixin):
    """
    Fish cages contain the fish.
    """

    def __init__(self, cage_id: int, cfg: Config, farm: Farm,
                 initial_lice_pop: Optional[GrossLiceDistrib] = None):
        """
        Create a cage on a farm.

        :param cage_id: the label (id) of the cage within the farm
        :param cfg: the farm configuration
        :param farm: a Farm object
        :param initial_lice_pop: if provided, overrides default generated lice population
        """
        super().__init__()

        self.cfg = cfg
        self.id = cage_id

        self.farm_id = farm.name
        self.start_date = cfg.farms[self.farm_id].cages_start[cage_id]
        self.date = cfg.start_date

        # TODO: update with calculations
        if initial_lice_pop is None:
            lice_population = {"L1": cfg.min_ext_pressure, "L2": 0, "L3": 0, "L4": 0, "L5f": 0,
                               "L5m": 0}
        else:
            lice_population = initial_lice_pop

        self.farm = farm

        self.egg_genotypes = GenoDistrib()
        # TODO/Question: what's the best way to deal with having multiple possible genetic schemes?
        # TODO/Question: I suppose some of this initial genotype information ought to come from the config file
        # TODO/Question: the genetic mechanism will be the same for all lice in a simulation, so should it live in the driver?
        self.genetic_mechanism = self.cfg.genetic_mechanism

        geno_by_lifestage = {stage: GenoDistrib(self.cfg.initial_genetic_ratios).normalise_to(lice_population[stage])
                             for stage in lice_population}

        self.lice_population = LicePopulation(geno_by_lifestage, self.cfg.initial_genetic_ratios)

        self.num_fish = cfg.farms[self.farm_id].num_fish
        self.num_infected_fish = self.get_mean_infected_fish()

        self.hatching_events: PriorityQueue[EggBatch] = PriorityQueue()
        self.arrival_events: PriorityQueue[TravellingEggBatch] = PriorityQueue()
        self.treatment_events: PriorityQueue[TreatmentEvent] = PriorityQueue()

        self.last_effective_treatment: Optional[TreatmentEvent] = None

    def to_json_dict(self):
        """
        Create a JSON-serialisable dictionary version of a cage.
        """
        filtered_vars = vars(self).copy()

        del filtered_vars["farm"]
        del filtered_vars["cfg"]
        del filtered_vars["logged_data"]

        # May want to improve these or change them if we change the representation for genotype distribs
        # Note: JSON encoders do not allow a default() being run on the _keys_ of dictionaries
        # and GenotypePopulation is a dict subclass. The simplest option here is to trivially force to_json_dict()
        # whenever possible,
        filtered_vars["egg_genotypes"] = self.egg_genotypes.to_json_dict()
        filtered_vars["lice_population"] = self.lice_population.to_json_dict()
        filtered_vars["geno_by_lifestage"] = self.lice_population.geno_by_lifestage.to_json_dict()
        filtered_vars["genetic_mechanism"] = str(filtered_vars["genetic_mechanism"])[len("GeneticMechanism."):]

        return filtered_vars

    def __str__(self):
        """
        Get a human readable string representation of the cage in json form.
        :return: a description of the cage
        """

        return json.dumps(self.to_json_dict(), cls=CustomFarmEncoder, indent=4)

    def update(
        self,
        cur_date: dt.datetime,
        pressure: int,
        ext_pressure_ratio: GenoDistribDict
    ) -> Tuple[GenoDistrib, Optional[dt.datetime], Money]:
        """Update the cage at the current time step.

        :param cur_date: Current date of simulation
        :param pressure: External pressure, planctonic lice coming from the reservoir
        :param ext_pressure_ratio: The genotype ratio to use for the external pressure
        :return: Tuple (egg genotype distribution, hatching date, cost)
        """

        self.clear_log()

        if cur_date >= self.start_date:
            logger.debug("\tUpdating farm {} / cage {}".format(self.farm_id, self.id))
            logger.debug("\t\tinitial fish population = {}".format(self.num_fish))
        else:
            logger.debug("\tUpdating farm {} / cage {} (non-operational)".format(self.farm_id, self.id))

        logger.debug("\t\tinitial lice population = {}".format(self.lice_population))

        logger.debug(f"\t\tAdding {pressure} lice from external pressure")

        # Background lice mortality events
        dead_lice_dist = self.get_background_lice_mortality()

        # Development events
        new_L2, new_L4, new_females, new_males = self.get_lice_lifestage(cur_date)

        # Lice coming from other cages and farms
        # NOTE: arrivals first then hatching
        hatched_arrivals_dist = self.get_arrivals(cur_date)

        # Egg hatching
        new_offspring_distrib = self.create_offspring(cur_date)

        # Lice coming from reservoir
        lice_from_reservoir = self.get_reservoir_lice(pressure, ext_pressure_ratio)
        logger.debug(f"\t\tExternal pressure lice distribution = {lice_from_reservoir}")

        if cur_date < self.start_date or self.is_fallowing:
            # Values that are not used before the start date
            treatment_mortality = self.lice_population.get_empty_geno_distrib()
            fish_deaths_natural, fish_deaths_from_lice = 0, 0
            num_infection_events = 0
            avail_dams_batch: OptionalDamBatch = None
            new_egg_batch: OptionalEggBatch = None
            cost = self.cfg.monthly_cost / 28

        else:
            # Events that happen when cage is populated with fish
            # (after start date)
            days_since_start = (cur_date - self.date).days

            # Treatment mortality events
            treatment_mortality, cost = self.get_lice_treatment_mortality(cur_date)

            # Fish growth and death
            fish_deaths_natural, fish_deaths_from_lice = self.get_fish_growth(days_since_start)

            # Infection events
            num_infection_events = self.do_infection_events(days_since_start)

            # Mating events that create eggs
            self.lice_population.free_dams(cur_date)
            delta_avail_dams, delta_eggs = self.do_mating_events(cur_date)
            avail_dams_batch = DamAvailabilityBatch(cur_date + dt.timedelta(days=self.cfg.dam_unavailability),
                                                    delta_avail_dams)

            new_egg_batch = self.get_egg_batch(cur_date, delta_eggs)

        fish_deaths_from_treatment = self.get_fish_treatment_mortality(
            (cur_date - self.start_date).days,
            fish_deaths_from_lice,
            fish_deaths_natural
        )

        self.update_deltas(
            dead_lice_dist,
            treatment_mortality,
            fish_deaths_natural,
            fish_deaths_from_lice,
            fish_deaths_from_treatment,
            new_L2,
            new_L4,
            new_females,
            new_males,
            num_infection_events,
            lice_from_reservoir,
            avail_dams_batch,
            new_offspring_distrib,
            hatched_arrivals_dist,
        )

        logger.debug("\t\tfinal lice population= {}".format(self.lice_population))

        egg_distrib = GenoDistrib()
        hatch_date: Optional[dt.datetime] = None
        if cur_date >= self.start_date and new_egg_batch:
            logger.debug("\t\tfinal fish population = {}".format(self.num_fish))
            egg_distrib = new_egg_batch.geno_distrib
            hatch_date = new_egg_batch.hatch_date
            logger.debug("\t\tlice offspring = {}".format(sum(egg_distrib.values())))

        assert egg_distrib.is_positive()
        return egg_distrib, hatch_date, cost

    def get_temperature(self, cur_date: dt.datetime) -> float:
        cur_month = cur_date.month
        return self.farm.year_temperatures[cur_month - 1]

    def get_lice_treatment_mortality_rate(self, cur_date: dt.datetime) -> GenoTreatmentDistrib:
        """Check if the cage is currently being treated. If yes, calculate the treatment rates.
        Note: this method consumes the internal treatment event queue.

        :param cur_date: the current date
        :returns: the mortality rates broken down by geno data.
        """
        susceptible_populations = [self.lice_population.geno_by_lifestage[stage] for stage in LicePopulation.susceptible_stages]
        num_susc_per_geno = cast(GenoDistrib, GenoDistrib.batch_sum(susceptible_populations))

        geno_treatment_distrib = {geno: GenoTreatmentValue(0.0, 0) for geno in num_susc_per_geno}

        def cts(event):
            nonlocal self
            self.last_effective_treatment = event

        pop_from_queue(self.treatment_events, cur_date, cts)

        # if no treatment has been applied check if the previous treatment is still effective
        if self.last_effective_treatment is None or \
                self.last_effective_treatment.affecting_date > cur_date or \
                cur_date > self.last_effective_treatment.treatment_window:
            return geno_treatment_distrib

        treatment_type = self.last_effective_treatment.treatment_type
        ave_temp = self.get_temperature(cur_date)

        logger.debug("\t\ttreating farm {}/cage {} on date {}".format(self.farm_id,
                                                                      self.id, cur_date))

        geno_treatment_distrib = self.cfg.get_treatment(treatment_type).get_lice_treatment_mortality_rate(self.lice_population, ave_temp)

        return geno_treatment_distrib

    def get_lice_treatment_mortality(self, cur_date) -> Tuple[GenoLifeStageDistrib, Money]:
        """
        Calculate the number of lice in each stage killed by treatment.

        Note: this method consumes the internal event queue

        :param cur_date: the current date
        """

        dead_lice_dist = self.lice_population.get_empty_geno_distrib()

        dead_mortality_distrib = self.get_lice_treatment_mortality_rate(cur_date)

        cost = Money("0.00")

        for geno, (mortality_rate, num_susc) in dead_mortality_distrib.items():
            if mortality_rate > 0:
                num_dead_lice = self.cfg.rng.poisson(mortality_rate * num_susc)
                num_dead_lice = min(num_dead_lice, num_susc)

                # We emulate the top algorithm with a multivariate hypergeom distrib
                population_by_stages = np.array([self.lice_population.geno_by_lifestage[stage][geno]
                                                 for stage in LicePopulation.susceptible_stages])

                # TODO: why is num_dead_lice not clamped?
                num_dead_lice = min(population_by_stages.sum(), num_dead_lice)
                dead_lice_nums = self.cfg.rng.multivariate_hypergeometric(
                    population_by_stages, num_dead_lice).tolist()
                for stage, dead_lice_num in zip(LicePopulation.susceptible_stages, dead_lice_nums):
                    dead_lice_dist[stage][geno] = dead_lice_num

                logger.debug("\t\tdistribution of dead lice on farm {}/cage {} = {}"
                             .format(self.farm_id, self.id, dead_lice_dist))

        # Compute cost
        if self.last_effective_treatment is not None and self.last_effective_treatment.end_application_date > cur_date:
            treatment_type = self.last_effective_treatment.treatment_type
            treatment_cfg = self.cfg.get_treatment(treatment_type)
            cage_days = (cur_date - self.start_date).days
            if isinstance(treatment_cfg, ChemicalTreatment):
                cost = treatment_cfg.price_per_kg * int(self.average_fish_mass(cage_days) / 1e3)
            elif isinstance(treatment_cfg, ThermalTreatment):
                cost = treatment_cfg.price_per_application

        return dead_lice_dist, cost

    def get_stage_ages_distrib(self, stage: str, temp_c: float = 10.0):
        """
        Create an age distribution (in days) for the sea lice within a lifecycle stage.

        This distribution is computed by using a simplified version of formulae (4), (6), (8)
        in Aldrin et al.: we assume that all lice of a stage m-1 that evolve at stage m
        will be put at a stage-age of 0, and that this is going to be a constant/stable
        amount.
        """

        stage_age_max_days = int(self.cfg.stage_age_evolutions[stage])

        if stage in ("L2", "L5f"):
            # Bogus uniform distribution L5m/L5f follow different schemes.
            # More realistically, this is an exponential/poisson distribution
            return np.full(stage_age_max_days, 1/stage_age_max_days)

        delta_p = self.cfg.delta_p[stage]
        delta_m10 = self.cfg.delta_m10[stage]
        delta_s = self.cfg.delta_s[stage]

        ages = np.arange(stage_age_max_days)

        weibull_median_rates = self._dev_rates(delta_p, delta_m10, delta_s, temp_c, ages)

        probas = np.empty_like(ages, dtype=np.float64)
        probas[0] = 1.0
        for i in range(1, stage_age_max_days):
            probas[i] = probas[i-1] * (1.0 - weibull_median_rates[i-1])

        return probas / np.sum(probas)

    @staticmethod
    def _dev_rates(del_p: float, del_m10: float, del_s: float, temp_c: float, ages: np.ndarray):
        """
        Probability of developing after n_days days in a stage as given by Aldrin et al 2017
        See section 2.2.4 of Aldrin et al. (2017)
        :param del_p: power transformation constant on temp_c
        :param del_m10: 10 °C median reference value
        :param del_s: fitted Weibull shape parameter
        :param temp_c: average temperature in °C
        :param ages: stage-age
        :return: expected development rate
        """
        epsilon = np.float64(1e-30)
        del_m = del_m10 * (10 / temp_c) ** del_p

        unbounded = np.clip(np.log(2) * del_s * ages ** (del_s - 1) * del_m ** (-del_s), epsilon, np.float64(1.0))
        return unbounded

    def get_lice_lifestage(self, cur_date: dt.datetime) -> Tuple[int, int, int, int]:
        """
        Move lice between lifecycle stages.
        See Section 2.1 of Aldrin et al. (2017)

        :param cur_date: the current date
        :returns: a tuple (new_l2, new_l4, new_l5f, new_l5m)
        """
        logger.debug("\t\tupdating lice lifecycle stages")

        def evolve_next_stage(stage: LifeStage, temp_c: float) -> int:
            # weibull params
            del_p = self.cfg.delta_p[stage]
            del_m10 = self.cfg.delta_m10[stage]
            del_s = self.cfg.delta_s[stage]
            num_lice = self.lice_population[stage]
            if num_lice == 0:
                return 0
            max_dev_time = self.cfg.stage_age_evolutions[stage]
            ages_distrib = self.get_stage_ages_distrib(stage, temp_c)
            # let us sample from 1000 lice
            ages = self.cfg.rng.choice(np.arange(max_dev_time), 1000, p=ages_distrib)

            evolution_rates = self._dev_rates(del_p, del_m10, del_s, temp_c, ages)
            average_rates = np.mean(evolution_rates)

            # print(f"Stage {stage} -> {average_rates}")

            return round(num_lice * average_rates) #int(min(self.cfg.rng.poisson(num_lice * average_rates), num_lice))

        lice_dist = {}
        ave_temp = self.get_temperature(cur_date)

        # L4 -> L5
        l4_to_l5 = evolve_next_stage("L4", ave_temp)
        new_females = int(self.cfg.rng.choice([math.floor(l4_to_l5 / 2.0), math.ceil(l4_to_l5 / 2.0)]))
        new_males = (l4_to_l5 - new_females)

        lice_dist["L5f"] = new_females
        lice_dist["L5m"] = new_males

        # L3 -> L4
        new_L4 = lice_dist["L4"] = evolve_next_stage("L3", ave_temp)

        # L2 -> L3 is done in do_infection_events(). Indeed, evolution from L2 to L3 is (virtually) age-independent

        # L1 -> L2
        new_L2 = lice_dist["L2"] = evolve_next_stage("L1", ave_temp)

        logger.debug("\t\t\tdistribution of new lice lifecycle stages on farm {}/cage {} = {}"
                     .format(self.farm_id, self.id, lice_dist))

        return new_L2, new_L4, new_females, new_males

    def get_fish_treatment_mortality(
            self,
            days_since_start: int,
            fish_lice_deaths: int,
            fish_backgroud_deaths: int) -> int:
        """
        Get fish mortality due to treatment. Mortality due to treatment is defined in terms of
        point percentage increases, thus we can only account for "excess deaths".
        If treatment is not being applied the overall

        :param days_since_start: the number of days since the beginning of the simulation
        :param fish_lice_deaths: the number of fish dead by lice
        :param fish_background_deaths: the number of fish dead by natural reasons

        :returns: number of fish death events
        """
        # See surveys/overton_treatment_mortalities.py for an explanation on what is going on
        mortality_events = fish_lice_deaths + fish_backgroud_deaths

        if mortality_events == 0 or self.num_fish == 0:
            return 0

        cur_date = self.start_date + dt.timedelta(days=days_since_start)
        if not self.is_treated(cur_date):
            return 0

        efficacy_window = self.last_effective_treatment.effectiveness_duration_days

        temperature = self.get_temperature(cur_date)
        fish_mass = self.average_fish_mass(days_since_start)

        last_treatment_params = self.cfg.get_treatment(self.last_effective_treatment.treatment_type)
        predicted_deaths = last_treatment_params.get_fish_mortality_occurrences(
            temperature, fish_mass, self.num_fish, efficacy_window, mortality_events)

        treatment_mortality_occurrences = round(predicted_deaths)

        return treatment_mortality_occurrences

    def get_fish_growth(self, days_since_start) -> Tuple[int, int]:
        """
        Get the number of fish that get killed either naturally or by lice.

        :param: days_since_start: the number of days since the beginning.
        :returns: a tuple (natural_deaths, lice_induced_deaths, lice_deaths)
        """

        logger.debug("\t\tupdating fish population")

        def fb_mort(days):
            """
            Fish background mortality rate, decreasing as in Soares et al 2011

            Fish death rate: constant background daily rate 0.00057 originally based on
            www.gov.scot/Resource/0052/00524803.pdf

            According to Volsett it should be higher (around 0.005) , but that's unlikely.
            However the authors used a trick: assuming that a mortality event takes 5 days one can
            divide the actual mortality rate by 5.
            Thus, 0.005 / 5 ~~ 0.0009-0.001 . It is still quite higher than expected.

            Therefore, we'll stick to this constant.
            TODO: what should we do with the formulae?

            :param days: number of days elapsed
            :return: fish background mortality rate
            """
            return 0.000057  # (1000 + (days - 700)**2)/490000000

        # Apply a sigmoid based on the number of lice per fish
        pathogenic_lice = sum([self.lice_population[stage] for stage in LicePopulation.pathogenic_stages])
        if self.num_infected_fish > 0:
            lice_per_host_mass = pathogenic_lice / (self.num_infected_fish *
                                                    (self.average_fish_mass(days_since_start) / 1e3))
        else:
            lice_per_host_mass = 0.0

        # The daily mortality rate also takes into account a mortality event takes days (in this case, 5).
        prob_lice_death = 0.2 / (1 + math.exp(-self.cfg.fish_mortality_k *
                                            (lice_per_host_mass - self.cfg.fish_mortality_center)))

        ebf_death = fb_mort(days_since_start) * self.num_fish
        elf_death = self.num_infected_fish * prob_lice_death
        fish_deaths_natural = round(ebf_death) # self.cfg.rng.poisson(ebf_death)
        fish_deaths_from_lice = round(elf_death) # self.cfg.rng.poisson(elf_death)

        logger.debug("\t\t\tnumber of background fish death {}, from lice {}"
                     .format(fish_deaths_natural, fish_deaths_from_lice))

        return fish_deaths_natural, fish_deaths_from_lice

    def compute_eta_aldrin(self, num_fish_in_farm, days):
        return self.cfg.infection_main_delta + math.log(num_fish_in_farm/1e5) + self.cfg.infection_weight_delta * \
               (math.log(self.average_fish_mass(days)/1e3) - self.cfg.delta_expectation_weight_log)

    def get_infection_rates(self, days_since_start) -> Tuple[float, int]:
        """
        Compute the number of lice that can infect and what their infection rate (number per fish) is

        :param days_since_start: days since the cage has opened

        :returns: a pair (Einf, num_avail_lice)
        """

        # Based on Aldrin et al.
        # Perhaps we can have a distribution which can change per day (the mean/median increaseѕ?
        # but at what point does the distribution mean decrease)./
        age_distrib = self.get_stage_ages_distrib("L2")
        num_avail_lice = round(self.lice_population["L2"] * np.sum(age_distrib[1:]))
        if num_avail_lice > 0:
            num_fish_in_farm = self.farm.num_fish

            # TODO: this has O(c^2) complexity
            etas = np.array([c.compute_eta_aldrin(num_fish_in_farm, days_since_start) for c in self.farm.cages])
            Einf = math.exp(etas[self.id]) / (1 + np.sum(np.exp(etas)))

            return Einf, num_avail_lice

        return 0.0, num_avail_lice

    def do_infection_events(self, days: int) -> int:
        """Infect fish in this cage if the sea lice are in stage L2 and at least 1 day old

        :param cur_date: current date of simulation
        :param days: days the number of days elapsed

        :return: number of evolving lice, or equivalently the new number of infections
        """
        Einf, num_avail_lice = self.get_infection_rates(days)

        if Einf == 0:
            return 0

        expected_events = Einf * num_avail_lice

        inf_events = self.cfg.rng.poisson(expected_events)

        return min(inf_events, num_avail_lice)

    def get_infecting_population(self, *args) -> int:
        if args is None or len(args) == 0:
            infective_stages = ["L3", "L4", "L5m", "L5f"]
        else:
            infective_stages = list(args)
        return sum(self.lice_population[stage] for stage in infective_stages)

    def get_mean_infected_fish(self, *args) -> int:
        """
        Get the average number of infected fish.

        :param \*args: the stages to consider (optional, by default all stages from the third onward are taken into account)

        :returns: the number of infected fish
        """
        if self.num_fish == 0:
            return 0

        attached_lice = self.get_infecting_population(*args)

        # see: https://stats.stackexchange.com/a/296053
        num_infected_fish = int(self.num_fish * (1 - ((self.num_fish - 1) / self.num_fish) ** attached_lice))
        return num_infected_fish

    def get_variance_infected_fish(self, n: int, k: int) -> float:
        # Rationale: assuming that we generate N bins that sum to K, we can model this as a multinomial distribution
        # where all p_i are the same. Therefore, the variance of each bin is k*(n-1)/(n**2)
        # Because we are considering the total variance of k events at the same time we need to multiply by k,
        # thus yielding k**2*(n-1)/(n**2).

        if n == 0:
            return 0.0
        return k**2 * (n - 1) / (n ** 2)

    def get_num_matings(self) -> int:
        """
        Get the number of matings. Implement Cox's approach assuming an unbiased sex distribution
        """

        # Background: AF and AM are randomly assigned to fish according to a negative multinomial distribution.
        # What we want is determining what is the expected likelihood there is _at least_ one AF and _at
        # least_ one AM on the same fish.

        males = self.lice_population["L5m"]
        females = self.lice_population.available_dams.gross

        if males == 0 or females == 0:
            return 0

        # VMR: variance-mean ratio; VMR = m/k + 1 -> k = m / (VMR - 1) with k being an "aggregation factor"
        mean = (males + females) / self.get_mean_infected_fish("L5m", "L5f")

        n = self.num_fish
        k = self.get_infecting_population("L5m", "L5f")
        variance = self.get_variance_infected_fish(n, k)
        vmr = variance / mean
        if vmr <= 1.0:
            return 0
        aggregation_factor = mean / (vmr - 1)

        prob_matching = 1 - (1 + males / ((males + females) * aggregation_factor)) ** (-1 - aggregation_factor)
        # TODO: using a poisson distribution can lead to a high std for high prob*females
        return int(np.clip(self.cfg.rng.poisson(prob_matching * females), np.int32(0), np.int32(min(males, females))))

    def do_mating_events(self, cur_date) -> Tuple[GenoDistrib, GenoDistrib]:
        """
        Will generate two deltas:  one to add to unavailable dams and subtract from available dams, one to add to eggs
        Assume males don't become unavailable? in this case we don't need a delta for sires

        :param cur_date: the current date

        :returns: a pair (mating_dams, new_eggs)
        """

        delta_eggs = GenoDistrib()
        num_matings = self.get_num_matings()

        distrib_sire_available = self.lice_population.geno_by_lifestage["L5m"]
        distrib_dam_available = self.lice_population.available_dams

        mating_dams = self.select_lice(distrib_dam_available, num_matings)
        mating_sires = self.select_lice(distrib_sire_available, num_matings)

        if distrib_sire_available.gross == 0 or distrib_dam_available.gross == 0:
            return GenoDistrib(), GenoDistrib()

        num_eggs = self.get_num_eggs(num_matings)
        if self.genetic_mechanism == GeneticMechanism.DISCRETE:
            delta_eggs = self.generate_eggs_discrete_batch(mating_sires, mating_dams, num_eggs)
        elif self.genetic_mechanism == GeneticMechanism.MATERNAL:
            delta_eggs = self.generate_eggs_maternal_batch(distrib_dam_available, num_eggs)

        delta_eggs = self.mutate(delta_eggs, mutation_rate=self.cfg.geno_mutation_rate)

        return mating_dams, delta_eggs

    def generate_eggs_discrete_batch(self, sire_distrib: GenoDistrib, dam_distrib: GenoDistrib,
                                     number_eggs: int) -> GenoDistrib:
        """
        Get number of eggs based on discrete genetic mechanism.

        The algorithm emulates the following scenario: each sire s_i is sampled from sire_distrib
        and will have a given genomic :math:`g_{s_i}` and similarly a dam :math:`d_i` will have a genotype :math:`g_{d_i}`.
        Because the genomic of a lice can only be dominant, recessive or partial dominant then
        the mating will result in one of these happening depending on the usual Mendelevian
        mechanism.
        The algorithm terminates when all sires and dams have been mated.
        Here we emulate such sampling strategy via a O(1) algorithm as follows: we compute
        the probabilities of each combination arising and adopt a multinomial distribution
        in order to achieve a given number of eggs.
        The rationale for the formulae used here is the following:

        - to get an A (dominant) one needs a combination of dominant and dominant allele, or dominant and the right half of partially dominant genes;

        - to get an a (recessive) one does the same as above, but with fully recessive alleles instead;

        - to get a partially dominant one, we use the inclusion-exclusion principle and subtract the cases above from all the possible pairings.

        Together with being fast, this approach naturally models statistical uncertainty compared
        to a perfect Mendelevian split. Unfortunately it does not naturally include mutation
        to non-existent strains (e.g. a mating between pure dominant lice distributions can never
        yield partial dominance or recessive genomics in the offspring).

        :param sire_distrib: the genotype distribution of the sires
        :param dam_distrib: the genotype distribution of the eggs
        :param number_eggs: the total number of eggs

        :returns: the newly sampled eggs as a :class:`GenoDistrib`.
        """

        assert sire_distrib.gross == dam_distrib.gross

        keys = [('A',), ('a',), ('A', 'a')]
        x1, y1, z1 = tuple(sire_distrib[k] for k in keys)
        x2, y2, z2 = tuple(dam_distrib[k] for k in keys)

        N_A = (x1 + 1/2*z1)*(x2+1/2*z2)
        N_a = (y1 + 1/2*z1)*(y2+1/2*z2)
        denom = sire_distrib.gross * dam_distrib.gross
        N_Aa = denom - N_A - N_a

        if denom == 0:
            return GenoDistrib()

        # We need to follow GenoDistrib's ordering now
        p = np.array([N_a, N_Aa, N_A]) / denom

        return GenoDistrib.from_ratios(number_eggs, p, self.cfg.rng)

    @staticmethod
    def generate_eggs_maternal_batch(dams: GenoDistrib, number_eggs: int) -> GenoDistrib:
        """Get number of eggs based on maternal genetic mechanism.

        Maternal-only inheritance - all eggs have mother's genotype.

        :param dam: the genomics of the dams
        :param number_eggs: the number of eggs produced
        :return: genomics distribution of eggs produced
        """
        return dams.normalise_to(number_eggs)

    def mutate(self, eggs: GenoDistrib, mutation_rate: float) -> GenoDistrib:
        """
        Mutate the genotype distribution

        :param eggs: the genotype distribution of the newly produced eggs
        :param mutation_rate: the rate of mutation with respect to the number of eggs.
        """
        if mutation_rate == 0:
            return eggs

        mutations = self.cfg.rng.poisson(mutation_rate * sum(eggs.values()))

        # generate a "swap" matrix
        # rationale: since ('a',) actually represents a pair of genes ('a', 'a')
        # there are only three directions: R->ID, ID->D, ID->R, D->ID. Note that
        # R->D or D->R are impossible via a single mutation.
        # Self-mutations are ignored.
        # To model this, we create a "masking" swapping matrix and force to 0 masked entries
        # and make sure they cannot be selected.

        alleles = [('a',), ('A',), ('A', 'a')]
        n = len(alleles)
        mask_matrix = np.array([[0, 0, 1],
                                [0, 0, 1],
                                [1, 1, 0]])

        p = mask_matrix.flatten() / np.sum(mask_matrix)

        # since I can't really be bothered to avoid negative mutations, we simply roll the dice every time until we get
        # a favourable outcome. Not ideal, but it works...

        while True:
            swap_matrix = self.cfg.rng.multinomial(mutations, p).reshape(n, n)
            result = eggs.copy()
            for idx, allele in enumerate(alleles):
                to_add = np.sum(swap_matrix[idx, :]) - np.sum(swap_matrix[:, idx])
                result[allele] += to_add
                if result[allele] == 0:
                    del result[allele]
            if result.is_positive():
                break

        return result

    def select_lice(self, distrib_lice_available: GenoDistrib, num_lice: int) -> GenoDistrib:
        """
        From a geno distribution of eligible lice sample a given Genotype distribution

        Note: this is very similar to GenoDistrib.normalise_to() but performs an explicit
        sampling.

        TODO: should we integrate this into GenoDistrib class

        :param distrib_lice_available: the starting dam genomic distribution
        :param num_lice: the wished number of dams to sample
        """
        # TODO: this function is flexible enough to be renamed and used elsewhere
        if sum(distrib_lice_available.values()) <= num_lice:
            return distrib_lice_available.copy()

        # "we need to select k lice from the given population broken down into different allele
        # bins and subtract" -> "select :math:`n` balls from the following :math`N_1, ..., N_k` bins without
        # replacement -> use a multivariate hypergeometric distribution
        lice_as_list = self.cfg.rng.multivariate_hypergeometric(
            list(distrib_lice_available.values()), num_lice)
        selected_lice = GenoDistrib(dict(zip(distrib_lice_available.keys(),
                                                   lice_as_list)))

        return selected_lice

    def get_num_eggs(self, mated_females) -> int:
        """
        Get the number of new eggs

        :param mated_females: the number of mated females that reproduce

        :returns: the number of eggs produced
        """

        # See Aldrin et al. 2017, §2.2.6
        age_distrib = self.get_stage_ages_distrib("L5f")
        age_range = np.arange(1, len(age_distrib) + 1)

        mated_females_distrib = mated_females * age_distrib

        # Hatching time is already covered in get_egg_batch
        eggs = self.cfg.reproduction_eggs_first_extruded * \
               (age_range ** self.cfg.reproduction_age_dependence) * mated_females_distrib

        return int(np.round(np.sum(eggs)))

    def get_egg_batch(self, cur_date: dt.datetime, egg_distrib: GenoDistrib) -> EggBatch:
        """
        Get the expected arrival date of an egg batch

        :param cur_date: the current time
        :param egg_distrib: the egg distribution

        :returns: EggBatch representing egg distribution and expected hatching date
        """

        # We use Stien et al (2005)'s regressed formula
        # τE= [β1/(T – 10 + β1β2)]**2  (see equation 8)
        # where β2**(-2) is the average temperature centered at around 10 degrees
        # and β1 is a shaping factor. This function is formally known as Belehrádek’s function
        cur_month = cur_date.month
        ave_temp = self.get_temperature(cur_date)

        beta_1 = self.cfg.delta_m10["L0_stien"]
        beta_2 = self.cfg.delta_p["L0"]
        expected_time = (beta_1 / (ave_temp - 10 + beta_1 * beta_2)) ** 2
        expected_hatching_date = cur_date + dt.timedelta(self.cfg.rng.poisson(expected_time))
        return EggBatch(expected_hatching_date, egg_distrib)

    def create_offspring(self, cur_time: dt.datetime) -> GenoDistrib:
        """
        Hatch the eggs from the event queue

        :param cur_time: the current time

        :returns: a delta egg genomics
        """

        delta_egg_offspring = GenoDistrib({geno: 0 for geno in self.cfg.initial_genetic_ratios})

        def cts(hatching_event: EggBatch):
            nonlocal delta_egg_offspring
            delta_egg_offspring += hatching_event.geno_distrib

        pop_from_queue(self.hatching_events, cur_time, cts)
        return delta_egg_offspring

    def get_arrivals(self, cur_date: dt.datetime) -> GenoDistrib:
        """Process the arrivals queue.

        :param cur_date: Current date of simulation

        :returns: Genotype distribution of eggs hatched in travel
        """

        hatched_dist = GenoDistrib()

        # check queue for arrivals at current date
        def cts(batch: TravellingEggBatch):
            nonlocal self
            nonlocal hatched_dist

            # if the hatch date is after current date, add to stationary egg queue
            if batch.hatch_date >= cur_date:
                stationary_batch = EggBatch(batch.hatch_date, batch.geno_distrib)
                self.hatching_events.put(stationary_batch)

            # otherwise the egg has hatched during travel, so determine the life stage
            # and add to population
            else:
                # TODO: determine life stage it arrives at; assumes all are L1 for now
                hatched_dist = batch.geno_distrib

        pop_from_queue(self.arrival_events, cur_date, cts)

        return hatched_dist

    def get_dying_lice_from_dead_fish(self, num_dead_fish: int) -> GrossLiceDistrib:
        """
        Get the number of lice that die when fish die.

        :param num_dead_fish: the number of dead fish
        :returns: a gross distribution
        """

        # Note: no paper actually provides a clear guidance on this.
        # This is mere speculation.
        # Basically, only consider PA and A males (for L4m = 0.5*L4)
        # as potentially able to survive and find a new host.
        # Only a fixed proportion can survive.

        if self.get_infecting_population() == 0 or self.num_infected_fish == 0:
            return {}

        affected_lice_gross = round(self.get_infecting_population() *
                                    num_dead_fish / self.num_infected_fish)
        infecting_lice = self.get_infecting_population()

        affected_lice_quotas = np.array([self.lice_population[stage] / infecting_lice * affected_lice_gross
                         for stage in LicePopulation.susceptible_stages])
        affected_lice_np = largest_remainder(affected_lice_quotas)

        affected_lice = Counter(dict(zip(LicePopulation.susceptible_stages, affected_lice_np.tolist())))

        surviving_lice_quotas = np.rint(np.trunc([
            self.lice_population['L4'] / (2 * infecting_lice) *
            affected_lice_gross * self.cfg.male_detachment_rate,
            self.lice_population['L5m'] / infecting_lice *
            affected_lice_gross * self.cfg.male_detachment_rate,
        ]))
        surviving_lice = Counter(dict(zip(['L4', 'L5m'], surviving_lice_quotas.tolist())))

        dying_lice = affected_lice - surviving_lice

        dying_lice_distrib = {k: int(v) for k, v in dying_lice.items() if v > 0}
        logger.debug(f"\t\tLice mortality due to fish mortality: {dying_lice_distrib}")
        return dying_lice_distrib

    def promote_population(
            self,
            prev_stage: Union[str, GenoDistrib],
            cur_stage: str,
            leaving_lice: int,
            entering_lice: Optional[int] = None
    ):
        """
        Promote the population by stage and respect the genotypes

        :param prev_stage: the lice stage from which cur_stage evolves
        :param cur_stage: the lice stage that is about to evolve
        :param leaving_lice: the number of lice in the _cur_stage=>next_stage_ progression
        :param entering_lice: the number of lice in the _prev_stage=>cur_stage_ progression. If _prev_stage_ is a a string, _entering_lice_ must be an _int_
        """
        if isinstance(prev_stage, str):
            if entering_lice is not None:
                prev_stage_geno = self.lice_population.geno_by_lifestage[prev_stage]
                entering_geno_distrib = prev_stage_geno.normalise_to(entering_lice)
            else:
                raise ValueError("entering_lice must be an int when prev_stage is a str")
        else:
            entering_geno_distrib = prev_stage
        cur_stage_geno = self.lice_population.geno_by_lifestage[cur_stage]

        leaving_geno_distrib = cur_stage_geno.normalise_to(leaving_lice)
        cur_stage_geno = cur_stage_geno + entering_geno_distrib - leaving_geno_distrib

        # update gross population. This is a bit hairy but I cannot think of anything simpler.
        self.lice_population.geno_by_lifestage[cur_stage] = cur_stage_geno

    def update_deltas(
            self,
            dead_lice_dist: GrossLiceDistrib,
            treatment_mortality: GenoLifeStageDistrib,
            fish_deaths_natural: int,
            fish_deaths_from_lice: int,
            fish_deaths_from_treatment: int,
            new_L2: int,
            new_L4: int,
            new_females: int,
            new_males: int,
            new_infections: int,
            lice_from_reservoir: Dict[LifeStage, GenoDistrib],
            delta_dams_batch: OptionalDamBatch,
            new_offspring_distrib: GenoDistrib,
            hatched_arrivals_dist: GenoDistrib,
    ):
        """Update the number of fish and the lice in each life stage

        :param dead_lice_dist: the number of dead lice due to background death (as a distribution)
        :param treatment_mortality: the distribution of genotypes being affected by treatment
        :param fish_deaths_natural: the number of natural fish death events
        :param fish_deaths_from_lice: the number of lice-induced fish death events
        :param fish_deaths_from_treatment: the number of treatment-induced fish death events
        :param new_L2: number of new L2 fish
        :param new_L4: number of new L4 fish
        :param new_females: number of new adult females
        :param new_males: number of new adult males
        :param new_infections: the number of new infections (i.e. progressions from L2 to L3)
        :param lice_from_reservoir: the number of lice taken from the reservoir
        :param delta_dams_batch: the genotypes of now-unavailable females in batch events
        :param new_offspring_distrib: the new offspring obtained from hatching and migrations
        :param hatched_arrivals_dist: new offspring obtained from arrivals
        """

        # Update dead_lice_dist to include fish-caused death as well
        dead_lice_by_fish_death = self.get_dying_lice_from_dead_fish(
            min(fish_deaths_natural + fish_deaths_from_lice, self.num_infected_fish))
        for affected_stage, reduction in dead_lice_by_fish_death.items():
            dead_lice_dist[affected_stage] += reduction

        for stage in self.lice_population:
            # update background mortality
            bg_delta = self.lice_population[stage] - dead_lice_dist[stage]
            self.lice_population[stage] = max(0, bg_delta)

            # update population due to treatment
            # TODO: __isub__ here is broken
            self.lice_population.geno_by_lifestage[stage] = self.lice_population.geno_by_lifestage[stage] - \
                                                            treatment_mortality[stage]

        self.lice_population.remove_negatives()

        self.promote_population("L4", "L5m", 0, new_males)
        self.promote_population("L4", "L5f", 0, new_females)
        self.promote_population("L3", "L4", new_males + new_females, new_L4)
        self.promote_population("L2", "L3", new_L4, new_infections)
        self.promote_population("L1", "L2", new_infections, new_L2)

        self.promote_population(new_offspring_distrib, "L1", new_L2, None)
        self.promote_population(hatched_arrivals_dist, "L1", 0, None)

        self.lice_population.remove_negatives()

        # in absence of wildlife genotype, simply upgrade accordingly
        self.lice_population.geno_by_lifestage["L2"] += lice_from_reservoir["L2"]
        self.lice_population.geno_by_lifestage["L1"] += lice_from_reservoir["L1"]

        if delta_dams_batch:
            self.lice_population.add_busy_dams_batch(delta_dams_batch)

        self.num_fish -= (fish_deaths_natural + fish_deaths_from_lice + fish_deaths_from_treatment)
        if self.num_fish < 0:
            self.num_fish = 0

        # treatment may kill some lice attached to the fish, thus update at the very end
        self.num_infected_fish = self.get_mean_infected_fish()

    # TODO: update arrivals dict type
    def update_arrivals(self, arrivals_dict: GenoDistribByHatchDate, arrival_date: dt.datetime):
        """Update the arrivals queue

        :param arrivals_dict: List of dictionaries of genotype distributions based on hatch date
        :param arrival_date: Arrival date at this cage
        """

        for hatch_date in arrivals_dict:

            # skip if there are no eggs in the dictionary
            if sum(arrivals_dict[hatch_date].values()) == 0:
                continue

            # create new travelling batch and update the queue
            batch = TravellingEggBatch(arrival_date, hatch_date, arrivals_dict[hatch_date])
            self.arrival_events.put(batch)

    def get_background_lice_mortality(self) -> GrossLiceDistrib:
        """
        Background death in a stage (remove entry) -> rate = number of
        individuals in stage*stage rate (nauplii 0.17/d, copepods 0.22/d,
        pre-adult female 0.05, pre-adult male ... Stien et al 2005)

        :returns: the current background mortality. The return value is genotype-agnostic

        """
        # TODO: this is dumb
        lice_mortality_rates = self.cfg.background_lice_mortality_rates
        lice_population = self.lice_population

        dead_lice_dist = {}
        for stage in lice_population:
            mortality_rate = lice_population[stage] * lice_mortality_rates[stage] # some exp / different mortalities - but this gets quickly unwieldly
            mortality: int = round(mortality_rate)#min(self.cfg.rng.poisson(mortality_rate), lice_population[stage])
            dead_lice_dist[stage] = mortality

        logger.debug("\t\tbackground mortality distribution of dead lice = {}".format(dead_lice_dist))
        return dead_lice_dist

    def average_fish_mass(self, days):
        """
        Average fish mass.

        :params days: number of elapsed days
        :returns: the average fish mass (in grams).
        """
        smolt_params = self.cfg.smolt_mass_params
        return smolt_params.max_mass / (1 + math.exp(smolt_params.skewness * (days - smolt_params.x_shift)))

    def get_reservoir_lice(self, pressure: int, external_pressure_ratios: GenoDistribDict) -> Dict[LifeStage, GenoDistrib]:
        """Get distribution of lice coming from the reservoir

        :param pressure: External pressure
        :param external_pressure_ratios: the external pressure ratios to sample from

        :return: Distribution of lice in L1 and L2
        """

        if pressure == 0:
            return {"L1": GenoDistrib(), "L2": GenoDistrib()}

        new_L1_gross = self.cfg.rng.integers(low=0, high=pressure, size=1)[0]
        new_L2_gross = pressure - new_L1_gross

        keys = list(external_pressure_ratios.keys())
        probas = list(external_pressure_ratios.values())
        new_L1 = GenoDistrib.from_ratios(new_L1_gross, probas, self.cfg.rng)
        new_L2 = GenoDistrib.from_ratios(new_L2_gross, probas, self.cfg.rng)
        #new_L1 = GenoDistrib(dict(zip(keys, self.cfg.rng.multinomial(new_L1_gross, probas).tolist())))
        #new_L2 = GenoDistrib(dict(zip(keys, self.cfg.rng.multinomial(new_L2_gross, probas).tolist())))

        new_lice_dist = {"L1": new_L1, "L2": new_L2}
        logger.debug("\t\tdistribution of new lice from reservoir = {}".format(new_lice_dist))
        return new_lice_dist

    def fallow(self):
        """Put the cage in a fallowing state.

        Implications:

        1. All the fish would be removed.
        2. L3/L4/L5 would therefore disappear as they are attached to fish.
        3. Dam waiting and treatment queues will be flushed.
        """

        for stage in LicePopulation.pathogenic_stages:
            self.lice_population[stage] = 0

        self.num_infected_fish = self.num_fish = 0
        self.treatment_events = PriorityQueue()

    @property
    def is_fallowing(self):
        return self.num_fish == 0

    def is_treated(self, cur_date):
        if self.last_effective_treatment is not None:
            treatment = self.last_effective_treatment
            if cur_date <= treatment.treatment_window:
                return True

        return False

    @property
    def aggregation_rate(self):
        """The aggregation rate is the number of lice over the total number of fish.
        Elsewhere it is referred to as infection rate, but here "infection rate" only refers to host fish.

        :returns: the aggregation rate"""
        return self.lice_population["L5f"] / self.num_fish if self.num_fish > 0 else 0.0


from __future__ import print_function, division
from clock_tree import ClockTree
import utils as ttutils
import config as ttconf
import numpy as np
from scipy import optimize as sciopt
from Bio import Phylo
from version import tt_version as __version__

class TreeTime(ClockTree):
    """
    TreeTime is a wrapper class to ClockTree that adds additional functionality
    such as reroot, detection and exclusion of outliers, resolution of polytomies
    using temporal information, and relaxed molecular clock models
    """

    def __init__(self, *args,**kwargs):
        """
        TreeTime constructor

        Parameters
        -----------
         Arguments to construct ClockTree

        Keyword Args
        ------------
         Kwargs to construct ClockTree

        """
        super(TreeTime, self).__init__(*args, **kwargs)
        self.n_iqd = ttconf.NIQD

    def run(self, root=None, infer_gtr=True, relaxed_clock=None, n_iqd = None,
            resolve_polytomies=True, max_iter=0, Tc=None, fixed_clock_rate=None,
            time_marginal=False, use_input_branch_length = False, **kwargs):

        """
        Run TreeTime reconstruction. Based on the input parameters, it divides
        the analysis into semi-independent jobs and conquers them one-by one
        gradually optimizing the tree given the temporal constarints and leaf
        nodes sequences.

        Parameters
        ----------

         root : str, None
            Try to find better root position on a given tree. If string is passed,
            the root will be searched according to the specified method. Available
            reroot methods are: 'best', 'oldest', '<leaf_name>'

            If None, use tree as-is.

         infer_gtr : bool default True
            Should infer GTR model?

         relaxed_clock : dic, None
            If not None, use autocorrelated molecular clock model. Specify the
            clock parameters as {slack:<slack>, coupling:<coupling>} dictionary.

         n_iqd : int, None
            If not None, filter tree nodes, which do not obey molecular clock
            for the particular tree. The nodes, which deviate more than
            :code:`n_iqd` interquantile intervals from the molecular clock
            regression will be marked as 'BAD' and not account in the TreeTime
            analysis

         resolve_polytomies : bool
            Should attempt to resolve multiple mergers?

         max_iter : int
            Maximum number of iterations to optimize the tree

         Tc : float, str, None
            If not None, use coalescent model to correct the branch lengths by
            introducing merger costs.

            If Tc is float, it is interpreted as the coalescence time scale

            If Tc is str, it should be one of (:code:`opt`, :code:`skyline`)

         fixed_clock_rate : float, None
            If None, infer clock rate from the molecular clock

            If float, use this rate

         time_marginal : bool default False
            Should perform marginal reconstruction of the node's positions?

         use_input_branch_length : bool
            If True, rely on the branch lengths in the imput tree and skip directly
            to the maximum-likelihood ancestral sequence reconstruction.
            Otherwise, perform preliminary sequence reconstruction using parsimony
            algorithm and do branch length optimization

        Keyword Args
        ------------

         Additional arguments needed by the dowstream funcitons

        """
        # determine how to reconstruct and sample sequences
        seq_kwargs = {"marginal":False, "sample_from_profile":"root"}
        if "fixed_pi" in kwargs:
            seq_kwargs["fixed_pi"] = kwargs["fixed_pi"]
        if "do_marginal" in kwargs:
            time_marginal=kwargs["do_marginal"]

        # initially, infer ancestral sequences and infer gtr model if desired
        if use_input_branch_length:
            self.infer_ancestral_sequences(infer_gtr=infer_gtr, **seq_kwargs)
            self.prune_short_branches()
        else:
            self.optimize_sequences_and_branch_length(infer_gtr=infer_gtr,
                                                  max_iter=2, prune_short=True, **seq_kwargs)
        avg_root_to_tip = np.mean([x.dist2root for x in self.tree.get_terminals()])

        # optionally reroot the tree either by oldest, best regression or with a specific leaf
        if n_iqd or root=='clock_filter':
            if "plot_rtt" in kwargs and kwargs["plot_rtt"]:
                plot_rtt=True
            else:
                plot_rtt=False
            self.clock_filter(reroot='best' if root=='clock_filter' else root,
                              n_iqd=n_iqd, plot=plot_rtt)
        elif root is not None:
            self.reroot(root=root)

        if use_input_branch_length:
            self.infer_ancestral_sequences(**seq_kwargs)
        else:
            self.optimize_sequences_and_branch_length(max_iter=2,**seq_kwargs)

        # infer time tree and optionally resolve polytomies
        self.logger("###TreeTime.run: INITIAL ROUND",0)
        self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=False, **kwargs)

        self.LH = [[self.tree.sequence_marginal_LH if seq_kwargs['marginal'] else self.tree.sequence_joint_LH,
                    self.tree.positional_joint_LH, 0.0]]

        # iteratively reconstruct ancestral sequences and re-infer
        # time tree to ensure convergence.
        niter = 0
        while niter < max_iter:

            self.logger("###TreeTime.run: ITERATION %d out of %d iterations"%(niter+1,max_iter),0)
            # add coalescent prior
            if Tc and (Tc is not None):
                from merger_models import Coalescent
                self.logger('TreeTime.run: adding coalescent prior with Tc='+str(Tc),1)
                self.merger_model = Coalescent(self.tree, Tc=avg_root_to_tip,
                                               date2dist=self.date2dist, logger=self.logger)

                if Tc=='skyline' and niter==max_iter-1: # restrict skyline model optimization to last iteration
                    self.merger_model.optimize_skyline(**kwargs)
                    self.logger("optimized a skyline ", 2)
                else:
                    if Tc in ['opt', 'skyline']:
                        self.merger_model.optimize_Tc()
                        self.logger("optimized Tc to %f"%self.merger_model.Tc.y[0], 2)
                    else:
                        try:
                            self.merger_model.set_Tc(Tc)
                        except:
                            self.logger("setting of coalescent time scale failed", 1, warn=True)

                self.merger_model.attach_to_tree()

            # estimate a relaxed molecular clock
            if relaxed_clock:
                self.relaxed_clock(**relaxed_clock)

            n_resolved=0
            if resolve_polytomies:
                # if polytomies are found, rerun the entire procedure
                n_resolved = self.resolve_polytomies()
                if n_resolved:
                    self.prepare_tree()
                    # when using the input branch length, only infer ancestral sequences
                    if use_input_branch_length:
                        self.infer_ancestral_sequences(**seq_kwargs)
                    else: # otherwise reoptimize branch length while preserving branches without mutations
                        self.optimize_sequences_and_branch_length(prune_short=False,
                                                                  max_iter=0, **seq_kwargs)

                    self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=False, **kwargs)
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                else:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                    self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=False, **kwargs)
            elif (Tc and (Tc is not None)) or relaxed_clock: # need new timetree first
                self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=False, **kwargs)
                ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
            else: # no refinements, just iterate
                ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=False, **kwargs)

            self.tree.coalescent_joint_LH = self.merger_model.total_LH() if Tc else 0.0

            self.LH.append([self.tree.sequence_marginal_LH if seq_kwargs['marginal'] else self.tree.sequence_joint_LH,
                            self.tree.positional_joint_LH, self.tree.coalescent_joint_LH])
            niter+=1

            if ndiff==0 & n_resolved==0:
                self.logger("###TreeTime.run: CONVERGED",0)
                break

        # if marginal reconstruction requested, make one more round with marginal=True
        # this will set marginal_pos_LH, which to be used as error bar estimations
        if time_marginal:
            self.logger("###TreeTime.run: FINAL ROUND - confidence estimation via marginal reconstruction", 0)
            self.make_time_tree(clock_rate=fixed_clock_rate, time_marginal=time_marginal, **kwargs)

    def clock_filter(self, reroot='best', n_iqd=None, plot=False):
        '''
        Labels outlier branches that don't seem to follow a molecular clock
        and excludes them from subsequent the molecular clock estimate and
        the timetree propagation

        Parameters
        ----------
         reroot : str, None
            Method to find the best root in the tree.

         n_iqd : int, None
            Number of iqd intervals. The outlier nodes are those which do not fall
            into :math:`IQD\cdot n_iqd` interval (:math:`IQD` is the interval between
            75 and 25 percentiles)

            if None, the default (3) assumed

         plot : bool
            Should plot the reults?

        '''
        if n_iqd is None:
            n_iqd = self.n_iqd

        terminals = self.tree.get_terminals()
        if reroot:
            self.reroot(root=reroot)
            icpt, clock_rate = self.tree.root._alpha, self.tree.root._beta
        else:
            tmp_date2dist = ttutils.DateConversion.from_tree(self.tree)
            icpt, clock_rate = tmp_date2dist.intercept, tmp_date2dist.clock_rate

        res = {}
        for node in terminals:
            if hasattr(node, 'numdate_given') and  (node.numdate_given is not None):
                res[node] = node.dist2root - clock_rate*np.mean(node.numdate_given) - icpt
        residuals = np.array(res.values())
        iqd = np.percentile(residuals,75) - np.percentile(residuals,25)
        for node,r in res.iteritems():
            if abs(r)>n_iqd*iqd and node.up.up is not None:
                self.logger('TreeTime.ClockFilter: marking %s as outlier, residual %f interquartile distances'%(node.name,r/iqd), 3)
                node.bad_branch=True
            else:
                node.bad_branch=False

        if plot:
            self.plot_root_to_tip()

        # redo root estimation after outlier removal
        if reroot:
            self.reroot(root=reroot)

    def plot_root_to_tip(self, add_internal=False, label=True, ax=None, **kwargs):
        """
        Plot root-to-tip regression

        Parameters
        ----------

         add_internal : bool
            Should plot internal node positoins?

         label : bool
            Should label the plots?

         ax: matplotlib axes, None
            If not None, use the provided matplotlib axes to plot the results

        Keyword Args
        ------------
         Additional arguments for matplotlib.pyplot.scatter function

        """
        import matplotlib.pyplot as plt
        tips = self.tree.get_terminals()
        internal = self.tree.get_nonterminals()
        if ax is None:
            plt.figure()
            ax=plt.subplot(111)
        dates = np.array([np.mean(n.numdate_given) for n in tips if n.numdate_given is not None])
        dist = np.array([n.dist2root for n in tips if n.numdate_given is not None])
        ind = np.array([n.bad_branch for n in tips if n.numdate_given is not None])
        # plot tips
        ax.scatter(dates[ind], dist[ind]  , c='r', label="bad tips" if label else "" , **kwargs)
        ax.scatter(dates[~ind], dist[~ind], c='g', label="good tips" if label else "", **kwargs)
        if add_internal and hasattr(self.tree.root, "numdate"):
            dates = np.array([n.numdate for n in internal])
            dist = np.array([n.dist2root for n in internal])
            ind = np.array([n.bad_branch for n in internal])
            # plot internal
            ax.scatter(dates[~ind], dist[~ind], c='b', marker='<', label="internal" if label else "", **kwargs)

        if label:
            ax.legend(loc=2)
        ax.set_ylabel('root-to-tip distance')
        ax.set_xlabel('date')
        ax.ticklabel_format(useOffset=False)
        plt.tight_layout()

    def reroot(self,root='best'):
        """
        Find best root and re-root the tree to the new root

        Parameters
        ----------

         root : str
            Which method should be used to find the best root. Available methods are:

            :code:`best` - maximize root-to-tip regression coefficient

            :code:`oldest` - choose the oldest node

            :code:`<node_name>` - reroot to the node with name :code:`<node_name>`

            :code:`[<node_name1>, <node_name2>, ...]` - reroot to the MRCA of these nodes
        """
        self.logger("TreeTime.reroot: with method or node: %s"%root,1)
        for n in self.tree.find_clades():
            n.branch_length=n.mutation_length
        from Bio import Phylo
        if isinstance(root,Phylo.BaseTree.Clade):
            new_root = root
        elif isinstance(root, list):
            new_root = self.tree.common_ancestor(*root)
        elif root in self._leaves_lookup:
            new_root = self._leaves_lookup[root]
        elif root=='oldest':
            new_root = sorted([n for n in self.tree.get_terminals()
                               if n.numdate_given is not None],
                               key=lambda x:np.mean(x.numdate_given))[0]
        elif root=='best':
            new_root = self.reroot_to_best_root(criterium='residual')
        elif root=='rsq':
            new_root = self.reroot_to_best_root(criterium='rsq')
        elif root=='residual':
            new_root = self.reroot_to_best_root(criterium='residual')
        elif root=='min_dev':
            new_root = self.reroot_to_best_root(criterium='min_dev')
        else:
            self.logger('TreeTime.reroot -- WARNING: unsupported rooting mechanisms or root not found',2,warn=True)
            return

        self.logger("TreeTime.reroot: Tree is being re-rooted to node "
                    +('new_node' if new_root.name is None else new_root.name), 2)
        if isinstance(root, list):
            #this forces a bifurcating root, as we want. Branch lengths will be reoptimized anyway.
            #(Without outgroup_branch_length, gives a trifurcating root, but this will mean
            #mutations may have to occur multiple times.)
            self.tree.root_with_outgroup(new_root, outgroup_branch_length=new_root.branch_length/2)
        else:
            self.tree.root_with_outgroup(new_root)
       # new nodes are produced when rooting with a terminal node, copy this clock info

        if new_root.is_terminal():
            if hasattr(new_root, "_alpha"):
                self.tree.root._alpha = new_root._alpha
            if hasattr(new_root, "_beta"):
                self.tree.root._beta = new_root._beta
            if hasattr(new_root, "_R2"):
                self.tree.root._R2 = new_root._R2
            if hasattr(new_root, "_residual"):
                self.tree.root._residual = new_root._residual

        self.tree.root.branch_length = self.one_mutation
        for n in self.tree.find_clades():
            n.mutation_length=n.branch_length
        self.tree.root.numdate_given = None
        # set root.gamma bc root doesn't have a branch_length_interpolator but gamma is needed
        if not hasattr(self.tree.root, 'gamma'):
            self.tree.root.gamma = 1.0
        self.prepare_tree()

    def resolve_polytomies(self, merge_compressed=False, rerun=True):
        """
        Resolve the polytomies on the tree.
        The function scans the tree, resolves polytomies in case there are any,
        and re-optimizes the tree with new topology. Note that polytomies are only
        resolved if that would result in higher likelihood. Sometimes, stretching
        two or more branches that carry several mutations are less costly than
        an additional branch with zero mutations (long branches are not stiff,
        short branches are.)

        Parameters
        ----------
         merge_compressed : bool
            Whether to keep compressed branches as polytomies or
            return a strictly binary tree.

        """
        self.logger("TreeTime.resolve_polytomies: resolving multiple mergers...",1)

        poly_found=0
        for n in self.tree.find_clades():
            if len(n.clades) > 2:
                prior_n_clades = len(n.clades)
                self._poly(n, merge_compressed)
                poly_found+=prior_n_clades - len(n.clades)

        obsolete_nodes = [n for n in self.tree.find_clades() if len(n.clades)==1 and n.up is not None]
        for node in obsolete_nodes:
            self.logger('TreeTime.resolve_polytomies: remove obsolete node '+node.name,4)
            if node.up is not None:
                self.tree.collapse(node)

        if poly_found:
            self.logger('TreeTime.resolve_polytomies: introduces %d new nodes'%poly_found,3)
        else:
            self.logger('TreeTime.resolve_polytomies: No more polytomies to resolve',3)
        return poly_found

    def _poly(self, clade, merge_compressed, verbose=1):

        """
        Function to resolve polytomies for a given parent node. If the
        number of the direct decendants is less than three (not a polytomy), does
        nothing. Otherwise, for each pair of nodes, assess the possible LH increase
        which could be gained by merging the two nodes. The increase in the LH is
        basically the tradeoff between the gain of the LH due to the changing the
        branch lenghts towards the optimal values and the decrease due to the
        introduction of the new branch with zero optimal length.
        """

        from branch_len_interpolator import BranchLenInterpolator
        from Bio import Phylo

        zero_branch_slope = self.gtr.mu*self.seq_len

        def _c_gain(t, n1, n2, parent):
            """
            cost gain if nodes n1, n2 are joined and their parent is placed at time t
            cost gain = (LH loss now) - (LH loss when placed at time t)
            """
            cg2 = n2.branch_length_interpolator(parent.time_before_present - n2.time_before_present) - n2.branch_length_interpolator(t - n2.time_before_present)
            cg1 = n1.branch_length_interpolator(parent.time_before_present - n1.time_before_present) - n1.branch_length_interpolator(t - n1.time_before_present)
            cg_new = - zero_branch_slope * (parent.time_before_present - t) # loss in LH due to the new branch
            return -(cg2+cg1+cg_new)

        def cost_gain(n1, n2, parent):
            """
            cost gained if the two nodes would have been connected.
            """
            try:
                cg = sciopt.minimize_scalar(_c_gain,
                    bounds=[max(n1.time_before_present,n2.time_before_present), parent.time_before_present],
                    method='Bounded',args=(n1,n2, parent))
                return cg['x'], - cg['fun']
            except:
                self.logger("TreeTime._poly.cost_gain: optimization of gain failed", 3, warn=True)
                return parent.time_before_present, 0.0

        def merge_nodes(source_arr, isall=False):
            mergers = np.array([[cost_gain(n1,n2, clade) if i1<i2 else (0.0,-1.0)
                                    for i1,n1 in enumerate(source_arr)]
                                for i2, n2 in enumerate(source_arr)])
            LH = 0
            while len(source_arr) > 1 + int(isall):
                # max possible gains of the cost when connecting the nodes:
                # this is only a rough approximation because it assumes the new node positions
                # to be optimal
                new_positions = mergers[:,:,0]
                cost_gains = mergers[:,:,1]
                # set zero to large negative value and find optimal pair
                np.fill_diagonal(cost_gains, -1e11)
                idxs = np.unravel_index(cost_gains.argmax(),cost_gains.shape)
                if (idxs[0] == idxs[1]) or cost_gains.max()<0:
                    self.logger("TreeTime._poly.merge_nodes: node is not fully resolved "+clade.name,4)
                    return LH

                n1, n2 = source_arr[idxs[0]], source_arr[idxs[1]]
                LH += cost_gains[idxs]

                new_node = Phylo.BaseTree.Clade()

                # fix positions and branch lengths
                new_node.time_before_present = new_positions[idxs]
                new_node.branch_length = clade.time_before_present - new_node.time_before_present
                new_node.clades = [n1,n2]
                n1.branch_length = new_node.time_before_present - n1.time_before_present
                n2.branch_length = new_node.time_before_present - n2.time_before_present

                # set parameters for the new node
                new_node.up = clade
                n1.up = new_node
                n2.up = new_node
                new_node.cseq = clade.cseq
                self._store_compressed_sequence_to_node(new_node)

                new_node.mutations = []
                new_node.mutation_length = 0.0
                new_node.branch_length_interpolator = BranchLenInterpolator(new_node, self.gtr, one_mutation=self.one_mutation)
                clade.clades.remove(n1)
                clade.clades.remove(n2)
                clade.clades.append(new_node)
                self.logger('TreeTime._poly.merge_nodes: creating new node as child of '+clade.name,3)
                self.logger("TreeTime._poly.merge_nodes: Delta-LH = " + str(cost_gains[idxs].round(3)), 3)

                # and modify source_arr array for the next loop
                if len(source_arr)>2: # if more than 3 nodes in polytomy, replace row/column
                    for ii in np.sort(idxs)[::-1]:
                        tmp_ind = np.arange(mergers.shape[0])!=ii
                        mergers = mergers[tmp_ind].swapaxes(0,1)
                        mergers = mergers[tmp_ind].swapaxes(0,1)

                    source_arr.remove(n1)
                    source_arr.remove(n2)
                    new_gains = np.array([[cost_gain(n1,new_node, clade) for n1 in source_arr]])
                    mergers = np.vstack((mergers, new_gains)).swapaxes(0,1)

                    source_arr.append(new_node)
                    new_gains = np.array([[cost_gain(n1,new_node, clade) for n1 in source_arr]])
                    mergers = np.vstack((mergers, new_gains)).swapaxes(0,1)
                else: # otherwise just recalculate matrix
                    source_arr.remove(n1)
                    source_arr.remove(n2)
                    source_arr.append(new_node)
                    mergers = np.array([[cost_gain(n1,n2, clade) for n1 in source_arr]
                                       for n2 in source_arr])

            return LH

        stretched = [c for c  in clade.clades if c.mutation_length < c.clock_length]
        compressed = [c for c in clade.clades if c not in stretched]

        if len(stretched)==1 and merge_compressed==False:
            return 0.0

        LH = merge_nodes(stretched, isall=len(stretched)==len(clade.clades))
        if merge_compressed and len(compressed)>1:
            LH += merge_nodes(compressed, isall=len(compressed)==len(clade.clades))

        return LH

    def print_lh(self, joint=True):
        """
        Print the total likelihood of the tree given the constrained leaves

        Parameters
        ----------

         joint : bool
            Whether joint or marginal LH should be printed

        """
        try:
            u_lh = self.tree.unconstrained_sequence_LH
            if joint:
                s_lh = self.tree.sequence_joint_LH
                t_lh = self.tree.positional_joint_LH
                c_lh = self.tree.coalescent_joint_LH
            else:
                s_lh = self.tree.sequence_marginal_LH
                t_lh = self.tree.positional_marginal_LH
                c_ls = 0

            print ("###  Tree Log-Likelihood  ###\n"
                " Sequence log-LH without constraints: \t%1.3f\n"
                " Sequence log-LH with constraints:    \t%1.3f\n"
                " TreeTime sequence log-LH:            \t%1.3f\n"
                " Coalescent log-LH:                   \t%1.3f\n"
               "#########################"%(u_lh, s_lh,t_lh, c_lh))
        except:
            print("ERROR. Did you run the corresponding inference (joint/marginal)?")

    def relaxed_clock(self, slack=None, coupling=None, **kwargs):
        """
        Allow the mutation rate to vary on the tree (relaxed molecular clock).
        Changes of the mutation rates from one branch to another are penalized.
        In addition, deviation of the mutation rate from the mean rate are
        penalized.

        Parameters
        ----------
         slack : float
            Maximum change in substitution rate between parent and child nodes

         coupling : float
            Maximum difference in substitution rates in sibling nodes

        """
        if slack is None: slack=ttconf.MU_ALPHA
        if coupling is None: coupling=ttconf.MU_BETA
        self.logger("TreeTime.relaxed_clock: slack=%f, coupling=%f"%(slack, coupling),2)

        c=1.0/self.one_mutation
        for node in self.tree.find_clades(order='postorder'):
            opt_len = node.mutation_length

            # opt_len \approx 1.0*len(node.mutations)/node.profile.shape[0] but calculated via gtr model
            # contact term: stiffness*(g*bl - bl_opt)^2 + slack(g-1)^2 =
            #               (slack+bl^2) g^2 - 2 (bl*bl_opt+1) g + C= k2 g^2 + k1 g + C
            node._k2 = slack + c*node.branch_length**2/(opt_len+self.one_mutation)
            node._k1 = -2*(c*node.branch_length*opt_len/(opt_len+self.one_mutation) + slack)
            # coupling term: \sum_c coupling*(g-g_c)^2 + Cost_c(g_c|g)
            # given g, g_c needs to be optimal-> 2*coupling*(g-g_c) = 2*child.k2 g_c  + child.k1
            # hence g_c = (coupling*g - 0.5*child.k1)/(coupling+child.k2)
            # substituting yields
            for child in node.clades:
                denom = coupling+child._k2
                node._k2 += coupling*(1.0-coupling/denom)**2 + child._k2*coupling**2/denom**2
                node._k1 += (coupling*(1.0-coupling/denom)*child._k1/denom \
                            - coupling*child._k1*child._k2/denom**2 \
                            + coupling*child._k1/denom)

        for node in self.tree.find_clades(order='preorder'):
            if node.up is None:
                node.gamma =- 0.5*node._k1/node._k2
            else:
                if node.up.up is None:
                    g_up = node.up.gamma
                else:
                    g_up = node.up.branch_length_interpolator.gamma
                node.branch_length_interpolator.gamma = (coupling*g_up - 0.5*node._k1)/(coupling+node._k2)

###############################################################################
### rerooting
###############################################################################
    def find_best_root_and_regression(self, criterium='rsq'):
        """
        Find the best root for the tree in linear time, given the timestamps of
        the leaves. The branch lengths should be optimized prior to the run;
        the terminal nodes should have the timestamps assigned as numdate_given
        attribute.
        """
        sum_ti =  np.sum([np.mean(node.numdate_given)*node.count for node in self.tree.get_terminals() if (not node.bad_branch)])
        sum_ti2 = np.sum([np.mean(node.numdate_given)**2*node.count for node in self.tree.get_terminals() if (not node.bad_branch)])
        N = 1.0*np.sum([node.count for node in self.tree.get_terminals() if not node.bad_branch])
        tip_count = 1.0*np.sum([1.0 for node in self.tree.get_terminals() if not node.bad_branch])
        if tip_count<2:
            self.logger("****ERROR: TreeTime.find_best_root_and_regression: need at least two dates to reroot!", 0, warn=True)
            self.logger("****ERROR: only %d tips have valid dates!"%N, 0, warn=True)
            return selt.tree.root, np.nan, np.nan

        Ninv = 1.0/N
        time_variance = (N*sum_ti2 - sum_ti**2)*Ninv**2

        #  fill regression terms for one of the two subtrees
        for node in self.tree.find_clades(order='postorder'):  # children first, msg to parents
            if node.is_terminal():  # inititalize the leaves
                #  will not rely on the standard func - count terminals directly
                node._st_n_leaves = 0 if node.bad_branch else node.count
                node._st_di = 0.0
                node._st_diti = 0.0
                node._st_di2 = 0.0

                if node.bad_branch:
                    node._st_ti = 0
                else:
                    node._st_ti = np.mean(node.numdate_given)*node.count

                node._ti = sum_ti
            else:
                #  for non-terminal nodes,
                node._st_ti = np.sum([k._st_ti for k in node.clades])
                node._st_n_leaves = np.sum([k._st_n_leaves for k in node.clades])
                node._st_di   = np.sum([k._st_di + k._st_n_leaves*k.branch_length for k in node.clades])
                node._st_diti = np.sum([k._st_diti + k.branch_length*k._st_ti for k in node.clades])
                node._st_di2  = np.sum([k._st_di2 + 2*k._st_di*k.branch_length + k._st_n_leaves*k.branch_length**2
                                       for k in node.clades])
                node._ti = sum_ti
                node.bad_branch = np.all([x.bad_branch for x in node])

        best_root = self.tree.root
        best_root_any = self.tree.root
        for node in self.tree.find_clades(order='preorder'):  # root first

            if node.up is None:
                # assign the values for the root node
                node._di   = node._st_di
                node._diti = node._st_diti
                node._di2  = node._st_di2

                dist_variance = (N*node._di2 - node._di**2)*(Ninv**2)
                disttime_cov = (N*node._diti - sum_ti*node._di)*(Ninv**2)
                time_variance = time_variance

                node._beta = disttime_cov/time_variance
                node._alpha = (node._di - node._beta*sum_ti)/N
                node._residual = (node._di2 - 2*node._beta*node._diti - 2*node._alpha*node._di
                                   + node._beta**2*sum_ti2 + 2*node._alpha*node._beta*sum_ti + node._alpha**2*N)

                node._R2 = disttime_cov**2/(time_variance*dist_variance)
                node._R2_delta_x = 0.0 # there is no branch to move the root
            elif node.bad_branch:
                node._beta = np.nan
                node._alpha = np.nan
                node._R2 = 0.0
                node._R2_delta_x = 0.0
                node._residual = np.inf

            elif criterium=='rsq': # calculate the r^2 of the root to tip regression and pick the best intermediate position on the branch
                #  NOTE order of these computation matters
                n_up = N - node._st_n_leaves
                n_down = node._st_n_leaves
                L = node.branch_length
                node._di = node.up._di + (n_up-n_down)*L
                node._di2 = (node.up._di2 + 2*L*node.up._di
                            - 4*(L*(node._st_di + n_down*L))
                            + N*L**2)
                node._diti = node.up._diti + L*(sum_ti - 2*node._st_ti)

                ## Express Node's sum_Di as the function of parent's sum_Di
                # and **displacement from parent's node x** :
                # sum_Di = A1 + A2 * x
                A1 = node.up._di
                A2 = n_up - n_down

                ## Express Node's sum_Di**2 as the function of parent's params
                # and **displacement from parent's node x** :
                # sum_Di2 = B1 + B2 * x + B3 * x**2
                B1 = node.up._di2
                B2 = 2 * (node.up._di - 2 * node._st_di - 2 * L * n_down )
                B3 = N

                ## Express Node's sum_DiTi as the function of parent's params
                # and **displacement from parent's node x** :
                # sum_DiTi = C1 + C2 * x
                C1 = node.up._diti
                C2 = sum_ti - 2 * node._st_ti

                ## Substituting Ai, Bi, Ci to the expression for R2, and
                ## making all the algebra, we get the R2 as the function of the
                ## displacement from the parent's node x:
                # R2(x) = CONST * (alpha * x**2 + beta * x+ gamma) / (mu * x**2 + nu * x + delta)
                # Expressions for alpha, beta, etc through Ai, Bi, Ci:
                alpha = (N * C2 - sum_ti * A2)**2
                beta = 2 * (N*C2 - sum_ti*A2) * (N*C1 - sum_ti*A1)
                gamma = (N*C1 - sum_ti*A1)**2
                mu = N * B3 - A2**2
                nu = N * B2 - 2 * A1 * A2
                delta = N * B1 - A1**2

                # Search for the maximum of R2 in the middle of the branch.
                # Eq: dR2/dx = 0 -> square equation:
                # x**2*(alpha*nu - beta *  mu) + 2x*(alpha*delta-mu*gamma) + (beta*delta - nu*gamma) = 0
                # look for the root(s):
                # Determinant is
                D2 =  (alpha * delta - mu * gamma) ** 2 - (alpha * nu - beta * mu) * (beta * delta - nu * gamma)

                if D2 < 0:
                    # somehow there is no extremum for the R2(x) function
                    x1 = -1.0 # any arbitrary value out of range [0, L], see below
                    x2 = -1.0
                else:
                    # actual roots - the extrema for the R2(x) function
                    if np.abs(alpha * nu - beta * mu)>0:
                        x1 = (-1 * (alpha * delta - mu * gamma) + D2**0.5) / (alpha * nu - beta * mu)
                        x2 = (-1 * (alpha * delta - mu * gamma) - D2**0.5) / (alpha * nu - beta * mu)
                    else:
                        x1 = -(beta*delta - nu*gamma)/(alpha * delta - mu * gamma)
                        x2 = x1

                # possible positions, where the new root can possibly be located
                # (restrict to the branch length)
                max_points = [k for k in (x1,x2,L) if k >= 0 and k <= L]
                # values of the R2 at these positions
                R2s = [(alpha * x**2 + beta * x + gamma) / (mu * x**2 + nu * x + delta) / time_variance / N**2 for x in max_points]
                # choose the best R2
                node._R2 = np.max(R2s)
                # and set the position for the best R2 value
                node._R2_delta_x = L - max_points[np.argmax(R2s)]

                # for this position, define the clock_rate and intercept:
                node._beta = ((L - node._R2_delta_x) * (N * C2 - sum_ti * A2) + (N*C1-sum_ti*A1)) / time_variance / N**2
                node._alpha = (L - node._R2_delta_x) * A2 / N  + (A1 - node._beta * sum_ti) / N
            elif criterium in ['residual', 'min_dev']: # calculate the squared residuals and minimize as rooting criterium
                L = node.branch_length
                # number of nodes descendent and outgrouping this node
                n_up = N - node._st_n_leaves
                n_down = node._st_n_leaves
                # sum of branch length of the descendent tree and the outgrouping one
                node._di = node.up._di + (n_up-n_down)*L
                nd_down = node._st_di
                nd_up = node._di - node._st_di
                # sum of times of descendent tips and outgrouping ones
                nt_down = node._st_ti
                nt_up = sum_ti - node._st_ti

                # sum of squared branch length of the descendent tree and the outgrouping one
                node._di2 = (node.up._di2 + 2*L*node.up._di
                            - 4*(L*(node._st_di + n_down*L))
                            + N*L**2)
                nd2_down = node._st_di2
                nd2_up = node._di2 - node._st_di2

                # sum of timexbranch length of the descendent tree and the outgrouping one
                node._diti = node.up._diti + L*(sum_ti - 2*node._st_ti)
                ndt_down = node._st_diti
                ndt_up = node._diti - node._st_diti

                disttime_cov = (N*node._diti - sum_ti*node._di)*(Ninv**2)

                # decompose expression for alpha and beta into parts independent and linear in epsilon
                # where epsilon is the shift in root position along the branch
                beta_0 = disttime_cov/time_variance
                b = L*(nt_down - nt_up - (n_down-n_up)*sum_ti*Ninv)*Ninv/time_variance
                alpha_0 = (node._di - beta_0*sum_ti)*Ninv
                a = (-b*sum_ti + L*(n_down-n_up))*Ninv
                eps = -(((nd_down - beta_0*nt_down - n_down*alpha_0) - (nd_up - beta_0*nt_up - n_up*alpha_0))/
                      ((n_down*L - b*nt_down  - a*n_down) - (-n_up*L - b*nt_up - a*n_up)))

                # only shifts between 0 and 1 are admissible (where 0 is the node itself and 1 is the parent)
                eps = min(1,max(0,eps))
                beta = beta_0 + eps*b
                alpha = alpha_0 + eps*a

                # calculate the residual and assign the regression coefficients
                node._residual = (node._di2 - 2*beta*node._diti - 2*alpha*node._di
                                   + beta**2*sum_ti2 + 2*alpha*beta*sum_ti + alpha**2*N)
                node._residual += N*(L*eps)**2
                node._residual += 2*eps*L*(nd_down - beta*nt_down - n_down*alpha) - 2*eps*L*(nd_up - beta*nt_up - n_up*alpha)
                node._alpha = alpha
                node._beta = beta
                node._R2_delta_x = eps*L
            else:
                self.logger("TreeTime.find_best_root_and_regression: unknown criterium",0)

            if criterium=='rsq':
                if node.up is None:
                    self.logger("TreeTime.find_best_root_and_regression: Initial root: R2:%f\tclock_rate:%.3e"%(best_root._R2, best_root._beta),3)
                if  node._R2 > best_root_any._R2:
                    best_root_any = node
                if  (node._R2 > best_root._R2 and node._beta>0) or best_root._beta<0:
                    best_root = node
                    self.logger("TreeTime.find_best_root_and_regression: Better root found: R2:%f\tclock_rate:%.3e\tbranch_displacement:%f"
                            %(best_root._R2, best_root._beta, (best_root._R2_delta_x) / ( best_root.branch_length + self.one_mutation)),4)
            elif criterium in ['residual', 'min_dev']:
                if node.up is None:
                    self.logger("TreeTime.find_best_root_and_regression: Initial root: residual:%.3e\tclock_rate:%.3e"%(best_root._residual, best_root._beta),3)
                if  node._residual < best_root_any._residual:
                    best_root_any = node
                if (node._residual < best_root._residual and node._beta>0) or best_root._beta<0:
                    best_root = node
                    self.logger("TreeTime.find_best_root_and_regression: Better root found: residual:%.3e\tclock_rate:%.3e\tbranch_displacement:%f"
                            %(best_root._residual, best_root._beta, (best_root._R2_delta_x) / ( best_root.branch_length + self.one_mutation)),4)

        if criterium=='rsq':
            if (best_root_any._R2 > best_root._R2):
                self.logger("WARNING: TreeTime.find_best_root_and_regression: optimal regression has negative rate: R2:%f\tclock_rate:%.3e"
                        %(best_root_any._R2, best_root_any._beta), 1)
            self.logger("TreeTime.find_best_root_and_regression: Best root: R2:%f\tclock_rate:%.3e\tbranch_displacement:%f"
                        %(best_root._R2, best_root._beta, (best_root._R2_delta_x) / ( best_root.branch_length + 0.001*self.one_mutation)),3)
        elif criterium=='residual':
            if (best_root_any._residual < best_root._residual):
                self.logger("WARNING: TreeTime.find_best_root_and_regression: optimal regression has negative rate: residual:%.3e\tclock_rate:%.3e"
                        %(best_root_any._residual, best_root_any._beta), 1)
            self.logger("TreeTime.find_best_root_an_R2_delta_xd_regression: Best root: residual:%.3e\tclock_rate:%.3e\tbranch_displacement:%f"
                        %(best_root._residual, best_root._beta, (best_root._R2_delta_x) / ( best_root.branch_length + 0.001*self.one_mutation)),3)
        elif criterium=='min_dev':
            if (best_root_any._residual < best_root._residual):
                self.logger("WARNING: TreeTime.find_best_root_and_regression: optimal regression has negative rate",1)
                best_root = best_root_any
            self.logger("TreeTime.find_best_root_an_R2_delta_xd_regression: Best root: residual:%.3e\tclock_rate:%.3e\tbranch_displacement:%f"
                        %(best_root._residual, best_root._beta, (best_root._R2_delta_x) / ( best_root.branch_length + 0.001*self.one_mutation)),3)

        return best_root, best_root._alpha, best_root._beta

    def reroot_to_best_root(self,infer_gtr = False, criterium='rsq', **kwarks):
        '''
        determine the node that, when the tree is rooted on this node, results
        in the best regression of temporal constraints and root to tip distances

        Parameters
        ----------

         infer_gtr : bool
            Should infer new GTR model after re-root?

        '''
        from Bio import Phylo
        self.logger("TreeTime.reroot_to_best_root: searching for the best root position...",2)
        best_root, a, b = self.find_best_root_and_regression(criterium=criterium)
        # first, re-root the tree

        if hasattr(best_root, "_R2_delta_x") and  best_root._R2_delta_x > 0 and best_root.up is not None:

            # create new node in the branch and root the tree to it
            new_node = Phylo.BaseTree.Clade()

            # insert the new node in the middle of the branch
            # by simple re-wiring the links on the both sides of the branch
            # and fix the branch lengths
            new_node.branch_length = best_root.branch_length - best_root._R2_delta_x
            new_node.up = best_root.up
            new_node._alpha = a
            new_node._beta = b
            new_node.clades = [best_root]
            if hasattr(best_root, "_R2"):
                new_node._R2 = best_root._R2
            if hasattr(best_root, "_residual"):
                new_node._residual = best_root._residual
            new_node.up.clades = [k if k != best_root else new_node
                                  for k in best_root.up.clades]

            best_root.branch_length = best_root._R2_delta_x
            best_root.up = new_node
            self.logger("TreeTime.reroot_to_best_root:"
                        " branch length of children of new root:"
                        " %.3e, %.3e"%(best_root.branch_length,
                                       new_node.branch_length),2)
            return new_node
        else:
            # simply use the existing node as the new root
            return best_root

def plot_vs_years(tt, years = 1, ax=None, confidence=None, ticks=True, **kwargs):
    '''
    converts branch length to years and plots the time tree on a time axis.
    Args:
        tt:     treetime object after a time tree is inferred
        years:  width of shaded boxes indicating blocks of years, default 1
        ax:     axis object. will create new axis of none specified
        confidence:     draw confidence intervals. This assumes that marginal
                        time tree inference was run
        **kwargs:   arbitrary kew word arguments that are passed down to Phylo.draw
    '''
    import matplotlib.pyplot as plt
    tt.branch_length_to_years()
    if ax is None:
        fig = plt.figure()
        ax = plt.subplot(111)
    # draw tree
    if "label_func" not in kwargs:
        nleafs = tt.tree.count_terminals()
        kwargs["label_func"] = lambda x:x.name if (x.is_terminal() and nleafs<30) else ""
    Phylo.draw(tt.tree, axes=ax, **kwargs)

    # set axis labels
    offset = tt.tree.root.numdate - tt.tree.root.branch_length
    xticks = ax.get_xticks()
    dtick = xticks[1]-xticks[0]
    shift = offset - dtick*(offset//dtick)
    tick_vals = [x+offset-shift for x in xticks]
    ax.set_xticks(xticks - shift)
    ax.set_xticklabels(map(str, tick_vals))
    ax.set_xlabel('year')
    ax.set_ylabel('')
    ax.set_xlim((0,np.max([n.numdate for n in tt.tree.get_terminals()])+2-offset))

    # put shaded boxes to delineate years
    if years:
        ylim = ax.get_ylim()
        xlim = ax.get_xlim()
        if type(years) in [int, float]:
            dyear=years
        from matplotlib.patches import Rectangle
        for yi,year in enumerate(np.arange(tick_vals[0], tick_vals[-1],dyear)):
            pos = year - offset
            r = Rectangle((pos, ylim[1]-5),
                          dyear, ylim[0]-ylim[1]+10,
                          facecolor=[0.7+0.1*(1+yi%2)] * 3,
                          edgecolor=[1,1,1])
            ax.add_patch(r)
            if year in tick_vals and pos>xlim[0] and pos<xlim[1] and ticks:
                ax.text(pos,ylim[0]-0.04*(ylim[1]-ylim[0]),str(int(year)),
                        horizontalalignment='center')
        ax.set_axis_off()

    # add confidence intervals to the tree graph -- grey bars
    if confidence:
        ttutils.tree_layout(tt.tree)
        if not hasattr(tt.tree.root, "marginal_inverse_cdf"):
            print("marginal time tree reconstruction required for confidence intervals")
        elif len(confidence)==2:
            cfunc = tt.get_confidence_interval
        elif len(confidence)==1:
            cfunc = tt.get_max_posterior_region
        else:
            print("confidence needs to be either a float (for max posterior region) or a two numbers specifying lower and upper bounds")
            return

        for n in tt.tree.find_clades():
            pos = cfunc(n, confidence)
            ax.plot(pos-offset, np.ones(len(pos))*n.ypos, lw=3, c=(0.5,0.5,0.5))

def treetime_to_newick(tt, outf):
    Phylo.write(tt.tree, outf, 'newick')

if __name__=="__main__":
    pass