from pprint import pprint

from birdgame.trackers.trackerbase import TrackerBase
from birdgame.stats.fewvar import FEWVar
import math
from densitypdf import density_pdf
import numpy as np
from birdgame.datasources.livedata import live_data_generator
from pprint import pprint


class MixtureTracker(TrackerBase):
    """
    A model that fits a mixture of two Gaussian distributions, one capturing the core
    distribution and another with a larger variance to capture the tails.
    """

    def __init__(self, fading_factor=0.0001, horizon=3):
        super().__init__(horizon)
        self.fading_factor = fading_factor
        self.current_x = None
        self.ewa_dx_core = FEWVar(fading_factor=fading_factor)
        self.ewa_dx_tail = FEWVar(fading_factor=fading_factor)
        self.weights = [0.9, 0.1]  # Heavily weight the core distribution

    def tick(self, payload):
        """
        Ingest a new record (payload), store it internally, and update the
        estimated Gaussian mixture model.
        """
        x = payload['dove_location']
        t = payload['time']
        self.add_to_quarantine(t, x)
        self.current_x = x
        prev_x = self.pop_from_quarantine(t)

        if prev_x is not None:
            x_change = x - prev_x

            # Winsorize the update for the core estimator to avoid tail effects
            threshold = 2.0 * math.sqrt(self.ewa_dx_core.get() if self.count > 0 else 1.0)
            if threshold > 0:
                winsorized_x_change = np.clip(x_change, -threshold, threshold)
            else:
                winsorized_x_change = x_change
            self.ewa_dx_core.update(winsorized_x_change)

            # Feed the tail estimator with double the real change magnitude
            self.ewa_dx_tail.update(2.0 * x_change)

            self.count += 1

    def predict(self):
        """
        Return a dictionary representing the best guess of the distribution,
        modeled as a mixture of two Gaussians.
        """
        x_mean = self.current_x
        components = []

        for i, ewa_dx in enumerate([self.ewa_dx_core, self.ewa_dx_tail]):
            try:
                x_var = ewa_dx.get()
                x_std = math.sqrt(x_var)
            except:
                x_std = 1.0

            if x_std <= 1e-6:
                x_std = 1e-6

            components.append({
                "density": {
                    "type": "builtin",
                    "name": "norm",
                    "params": {"loc": x_mean, "scale": x_std}
                },
                "weight": self.weights[i]
            })

        prediction_rec = {
            "type": "mixture",
            "components": components
        }

        # Verify the generated density using density_pdf
        _ = density_pdf(prediction_rec, x=0.0)  # Ensure the format is valid

        return prediction_rec


def example_of_testing_manually():
    # Just an example
    gen = live_data_generator()
    pdf = None
    for payload in gen:
        tracker.tick(payload)
        pdf = tracker.predict()
        if tracker.count > 100:
            break
    pprint(pdf)


if __name__ == '__main__':
    tracker = MixtureTracker()
    tracker.test_run(
        live=True, # Set to True to use live streaming data; set to False to use data from a CSV file
        step_print=10000 # Print the score and progress every 1000 steps
    )