import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from scipy.stats import pearsonr
import numpy as np
from scipy.fft import fft


class EventManager:
    def __init__(self):
        self.events = []

    def add_event(self, event_title, time_dataset, probability_fork, quantity, common_tag_event_dataset,
                  quantity_correlation_dataset, event_max_quantity, event_min_quantity, event_middle_quantity,
                  sentiment_direction):
        event = {
            "event_title": event_title,
            "time_dataset": time_dataset,
            "probability_fork": probability_fork,
            "quantity": quantity,
            "common_tag_event_dataset": common_tag_event_dataset,
            "quantity_correlation_dataset": quantity_correlation_dataset,
            "event_max_quantity": event_max_quantity,
            "event_min_quantity": event_min_quantity,
            "event_middle_quantity": event_middle_quantity,
            "sentiment_direction": sentiment_direction
        }
        self.events.append(event)

    def remove_event(self, event_title):
        self.events = [event for event in self.events if event['event_title'] != event_title]

    def get_events_by_tag(self, tag):
        return [event for event in self.events if tag in event['common_tag_event_dataset']]

    def get_events_by_sentiment(self, sentiment):
        return [event for event in self.events if event['sentiment_direction'] == sentiment]

    def get_events_by_quantity_range(self, min_quantity, max_quantity):
        return [event for event in self.events if min_quantity <= event['quantity'] <= max_quantity]

    def predict_time_series(self, event_title):
        event = next((event for event in self.events if event['event_title'] == event_title), None)
        if event:
            time_series = event['time_dataset']
            # Aquí puedes implementar tu modelo de predicción de series temporales
            # Por ejemplo, utilizando un modelo de regresión como RandomForestRegressor de scikit-learn
            X = np.arange(len(time_series)).reshape(-1, 1)
            y = np.array(time_series)
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
            model = RandomForestRegressor()
            model.fit(X_train, y_train)
            predictions = model.predict(X_test)
            return predictions
        else:
            return None

    def plot_event_parameters_over_time(self, event_title):
        event = next((event for event in self.events if event['event_title'] == event_title), None)
        if event:
            time_series = event['time_dataset']
            plt.plot(time_series)
            plt.xlabel('Tiempo')
            plt.ylabel('Valor')
            plt.title('Parámetros del Evento "{}" a lo largo del tiempo'.format(event_title))
            plt.show()

    def plot_prediction(self, event_title):
        predictions = self.predict_time_series(event_title)
        if predictions:
            plt.plot(predictions, label='Predicción')
            plt.xlabel('Tiempo')
            plt.ylabel('Valor')
            plt.title('Predicción del Evento "{}"'.format(event_title))
            plt.legend()
            plt.show()

    def check_correlation(self, event_title1, event_title2):
        event1 = next((event for event in self.events if event['event_title'] == event_title1), None)
        event2 = next((event for event in self.events if event['event_title'] == event_title2), None)
        if event1 and event2:
            correlation, _ = pearsonr(event1['quantity_correlation_dataset'], event2['quantity_correlation_dataset'])
            return correlation
        else:
            return None

    def fourier_transform(self, event_title):
        event = next((event for event in self.events if event['event_title'] == event_title), None)
        if event:
            time_series = event['time_dataset']
            transformed_data = fft(time_series)
            return transformed_data
        else:
            return None


# Ejemplo de uso
event_manager = EventManager()

# Añadir eventos
event_manager.add_event("Evento 1", [1, 2, 3, 4, 5], 0.8, 100, ["tag1", "tag2"], [0.1, 0.2, 0.3, 0.4, 0.5],
                        150, 50, 100, "good when up")
event_manager.add_event("Evento 2", [2, 4, 6, 8, 10], 0.6, 200, ["tag2", "tag3"], [0.2, 0.4, 0.6, 0.8, 1.0],
                        250, 150, 200, "bad when down")

# Realizar predicción de series temporales y plot
event_manager.plot_event_parameters_over_time("Evento 1")
event_manager.plot_prediction("Evento 1")

# Comprobar correlación entre dos eventos
correlation = event_manager.check_correlation("Evento 1", "Evento 2")
if correlation:
    print("Correlación entre Evento 1 y Evento 2:", correlation)
else:
    print("Alguno de los eventos no existe.")

# Transformada de Fourier
transformed_data = event_manager.fourier_transform("Evento 1")
print("Transformada de Fourier del Evento 1:", transformed_data)