import torch
from utils.autoregressive import predict_auto_regressive_quantile
from scipy.interpolate import interp1d
from trainers.trainer import Trainer
from trainers.autoregressive_trainer import AutoRegressiveTrainer
from data.preprocessing import DataProcessor
from utils.clearml import ClearMLHelper
from losses import PinballLoss, NonAutoRegressivePinballLoss
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt


class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
    def __init__(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer, data_processor: DataProcessor, quantiles: list, device: torch.device, clearml_helper: ClearMLHelper = None, debug: bool = True):
        quantiles_tensor = torch.tensor(quantiles)
        quantiles_tensor = quantiles_tensor.to(device)
        self.quantiles = quantiles

        criterion = PinballLoss(quantiles=quantiles_tensor)
        super().__init__(model=model, optimizer=optimizer, criterion=criterion, data_processor=data_processor, device=device, clearml_helper=clearml_helper, debug=debug)

    
    def log_final_metrics(self, task, dataloader, train: bool = True):
        metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }
        transformed_metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }

        with torch.no_grad():
            total_amount_samples = len(dataloader.dataset) - 95

            for idx in tqdm(range(total_amount_samples)):
                _, outputs, samples, targets = self.auto_regressive(dataloader, idx)


                inversed_samples = torch.tensor(self.data_processor.inverse_transform(samples))
                inversed_targets = torch.tensor(self.data_processor.inverse_transform(targets))

                outputs = outputs.to(self.device)
                targets = targets.to(self.device)
                samples = samples.to(self.device)

                for metric in self.metrics_to_track:
                    if metric.__class__ != PinballLoss:
                        transformed_metrics[metric.__class__.__name__] += metric(samples, targets)
                        metrics[metric.__class__.__name__] += metric(inversed_samples, inversed_targets)
                    else:
                        transformed_metrics[metric.__class__.__name__] += metric(outputs, targets)

            for metric in self.metrics_to_track:
                metrics[metric.__class__.__name__] /= total_amount_samples
                transformed_metrics[metric.__class__.__name__] /= total_amount_samples

            for metric_name, metric_value in metrics.items():
                if PinballLoss.__name__ in metric_name:
                    continue
                name = f'train_{metric_name}' if train else f'test_{metric_name}'
                task.get_logger().report_single_value(name=name, value=metric_value)

            for metric_name, metric_value in transformed_metrics.items():
                name = f'train_transformed_{metric_name}' if train else f'test_transformed_{metric_name}'
                task.get_logger().report_single_value(name=name, value=metric_value)

    def get_plot(self, current_day, next_day, predictions, show_legend: bool = True):
        fig = go.Figure()

        # Convert to numpy for plotting
        current_day_np = current_day.view(-1).cpu().numpy()
        next_day_np = next_day.view(-1).cpu().numpy()
        predictions_np = predictions.cpu().numpy()

        # Add traces for current and next day
        fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
        fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))

        for i, q in enumerate(self.quantiles):
            fig.add_trace(go.Scatter(x=96 + np.arange(96), y=predictions_np[:, i], 
                                    name=f"Prediction (Q={q})", line=dict(dash='dash')))

        # Update the layout
        fig.update_layout(title="Predictions and Quantiles of the Linear Model", showlegend=show_legend)
        
        return fig

    def auto_regressive(self, data_loader, idx, sequence_length: int = 96):
        self.model.eval()
        target_full = []
        predictions_sampled = []
        predictions_full = []

        prev_features, target = data_loader.dataset[idx]
        prev_features = prev_features.to(self.device)

        initial_sequence = prev_features[:96]

        target_full.append(target)
        with torch.no_grad():
            prediction = self.model(prev_features.unsqueeze(0))
        predictions_full.append(prediction.squeeze(0))

        # sample from the distribution
        sample = self.sample_from_dist(self.quantiles.cpu(), prediction.squeeze(-1).cpu().numpy())
        predictions_sampled.append(sample)

        for i in range(sequence_length - 1):
            new_features = torch.cat((prev_features[1:97].cpu(), torch.tensor([predictions_sampled[-1]])), dim=0)
            new_features = new_features.float()

            # get the other needed features
            other_features, new_target = data_loader.dataset.random_day_autoregressive(idx + i + 1)

            if other_features is not None:
                prev_features = torch.cat((new_features, other_features), dim=0)
            else:
                prev_features = new_features

            # add target to target_full
            target_full.append(new_target)

            # predict
            with torch.no_grad():
                prediction = self.model(new_features.unsqueeze(0).to(self.device))
            predictions_full.append(prediction.squeeze(0))

            # sample from the distribution
            sample = self.sample_from_dist(self.quantiles.cpu(), prediction.squeeze(-1).cpu().numpy())
            predictions_sampled.append(sample)

        return initial_sequence.cpu(), torch.stack(predictions_full).cpu(), torch.tensor(predictions_sampled).reshape(-1, 1), torch.stack(target_full).cpu()
    
    @staticmethod
    def sample_from_dist(quantiles, output_values):
        # Interpolate the inverse CDF
        inverse_cdf = interp1d(quantiles, output_values, kind='linear', bounds_error=False, fill_value="extrapolate")

        # generate one random uniform number
        uniform_random_numbers = np.random.uniform(0, 1, 1000)

        # Apply the inverse CDF to the uniform random numbers
        samples = inverse_cdf(uniform_random_numbers)

        # Return the mean of the samples
        return np.mean(samples)

    def plot_quantile_percentages(self, task, data_loader, train: bool = True, iteration: int = None):

        total = 0
        quantile_counter = {q: 0 for q in self.quantiles.cpu().numpy()}

        with torch.no_grad():
            for inputs, targets in data_loader:
                inputs = inputs.to("cuda")
                output = self.model(inputs)

                # output shape: (batch_size, num_quantiles)
                # target shape: (batch_size, 1)
                for i, q in enumerate(self.quantiles.cpu().numpy()):
                    quantile_counter[q] += np.sum(targets.squeeze(-1).cpu().numpy() < output[:, i].cpu().numpy())

                total += len(targets)

        # to numpy array of length len(quantiles)
        percentages = np.array([quantile_counter[q] / total for q in self.quantiles.cpu().numpy()])

        bar_width = 0.35
        index = np.arange(len(self.quantiles.cpu().numpy()))

        # Plotting the bars
        fig, ax = plt.subplots(figsize=(15, 10))

        bar1 = ax.bar(index, self.quantiles.cpu().numpy(), bar_width, label='Ideal', color='brown')
        bar2 = ax.bar(index + bar_width, percentages, bar_width, label='NN model', color='blue')

        # Adding the percentage values above the bars for bar2
        for rect in bar2:
            height = rect.get_height()
            ax.text(rect.get_x() + rect.get_width() / 2., 1.005 * height,
                    f'{height:.2}', ha='center', va='bottom')  # Format the number as a percentage

        series_name = "Training Set" if train else "Test Set"

        # Adding labels and title
        ax.set_xlabel('Quantile')
        ax.set_ylabel('Fraction of data under quantile forecast')
        ax.set_title(f'Quantile Performance Comparison ({series_name})')
        ax.set_xticks(index + bar_width / 2)
        ax.set_xticklabels(self.quantiles.cpu().numpy())
        ax.legend()

        task.get_logger().report_matplotlib_figure(title='Quantile Performance Comparison', series=series_name, report_image=True, figure=plt, iteration=iteration)
        plt.close()
                
class NonAutoRegressiveQuantileRegression(Trainer):
    def __init__(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer, data_processor: DataProcessor, quantiles: list, device: torch.device, clearml_helper: ClearMLHelper = None, debug: bool = True):
        quantiles_tensor = torch.tensor(quantiles)
        quantiles_tensor = quantiles_tensor.to(device)
        self.quantiles = quantiles

        criterion = NonAutoRegressivePinballLoss(quantiles=quantiles_tensor)
        super().__init__(model=model, optimizer=optimizer, criterion=criterion, data_processor=data_processor, device=device, clearml_helper=clearml_helper, debug=debug)

    @staticmethod
    def sample_from_dist(quantiles, output_values):
        reshaped_values = output_values.reshape(-1, len(quantiles))
        samples = []
        for row in reshaped_values:
            inverse_cdf = interp1d(quantiles, row, kind='linear', bounds_error=False, fill_value="extrapolate")
            uniform_random_numbers = np.random.uniform(0, 1, 1000)
            new_samples = inverse_cdf(uniform_random_numbers)
            samples.append(np.mean(new_samples))
        return np.array(samples)
        
    def log_final_metrics(self, task, dataloader, train: bool = True):
        metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }
        transformed_metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }

        with torch.no_grad():
            for inputs, targets in dataloader:
                inputs, targets = inputs.to(self.device), targets

                
                outputs = self.model(inputs)
                outputted_samples = [self.sample_from_dist(self.quantiles.cpu(), output.cpu().numpy()) for output in outputs]
                # to tensor
                outputted_samples = torch.tensor(outputted_samples)

                inversed_outputs = torch.tensor(self.data_processor.inverse_transform(outputted_samples))
                inversed_inputs = torch.tensor(self.data_processor.inverse_transform(targets))

                # set on same device
                inversed_outputs = inversed_outputs.to(self.device)
                inversed_inputs = inversed_inputs.to(self.device)
                outputted_samples = outputted_samples.to(self.device)

                for metric in self.metrics_to_track:
                    transformed_metrics[metric.__class__.__name__] += metric(outputted_samples, targets.to(self.device))
                    metrics[metric.__class__.__name__] += metric(inversed_outputs, inversed_inputs)

            for metric in self.metrics_to_track:
                metrics[metric.__class__.__name__] /= len(dataloader)
                transformed_metrics[metric.__class__.__name__] /= len(dataloader)

            for metric_name, metric_value in metrics.items():
                if train:
                    metric_name = f'train_{metric_name}'
                else:
                    metric_name = f'test_{metric_name}'
                
                task.get_logger().report_single_value(name=metric_name, value=metric_value)

            for metric_name, metric_value in transformed_metrics.items():
                if train:
                    metric_name = f'train_transformed_{metric_name}'
                else:
                    metric_name = f'test_transformed_{metric_name}'

                task.get_logger().report_single_value(name=metric_name, value=metric_value)

    
    def get_plot(self, current_day, next_day, predictions, show_legend: bool = True):
        fig = go.Figure()

        # Convert to numpy for plotting
        current_day_np = current_day.view(-1).cpu().numpy()
        next_day_np = next_day.view(-1).cpu().numpy()
        
        # reshape predictions to (n, len(quantiles))$
        predictions_np = predictions.cpu().numpy().reshape(-1, len(self.quantiles))

        # Add traces for current and next day
        fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
        fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))

        for i, q in enumerate(self.quantiles):
            fig.add_trace(go.Scatter(x=96 + np.arange(96), y=predictions_np[:, i], 
                                    name=f"Prediction (Q={q})", line=dict(dash='dash')))

        # Update the layout
        fig.update_layout(title="Predictions and Quantiles", showlegend=show_legend)
        
        return fig

