#!/usr/bin/env python3

import json
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from pathlib import Path
import glob
import argparse

class ConstraintOptimizationPlotter:
    def __init__(self, config_path):
        with open(config_path, 'r') as f:
            self.config = json.load(f)

        self.data_path = Path(self.config['data_path'])
        self.output_dir = Path(self.config['output_dir'])
        self.output_dir.mkdir(exist_ok=True)

        # Load objectives for percentage deviation calculation
        objectives_path = Path(__file__).parent / 'objectives.json'
        with open(objectives_path, 'r') as f:
            self.objectives = json.load(f)

    def get_csv_pattern(self, plot_type):
        patterns = {
            'best_candidates': 'best_candidates_*.csv',
            'constraint_violation': 'constraint_violation/constraint_violations_*.csv',
            'feasible_fraction': 'feasible_fraction/feasible_fractions_*.csv'
        }
        return patterns.get(plot_type, 'best_candidates_*.csv')

    def get_value_column(self, plot_type):
        columns = {
            'best_candidates': 'evaluation',  # Changed from 'fitness' to 'evaluation'
            'constraint_violation': 'avg_constraint_violation',
            'feasible_fraction': 'feasible_fraction'
        }
        return columns.get(plot_type, 'evaluation')

    def blend_colors(self, color1, color2, alpha=0.5):
        """Blend two colors together"""
        c1 = mcolors.to_rgb(color1)
        c2 = mcolors.to_rgb(color2)
        blended = tuple(alpha * c1[i] + (1 - alpha) * c2[i] for i in range(3))
        return mcolors.to_hex(blended)

    def get_combined_color(self, algorithm_color, instance_color, combination_method="blend"):
        """Combine algorithm and instance colors based on the specified method"""
        if combination_method == "blend":
            return self.blend_colors(algorithm_color, instance_color, alpha=0.6)
        elif combination_method == "algorithm":
            return algorithm_color
        elif combination_method == "instance":
            return instance_color
        else:
            return self.blend_colors(algorithm_color, instance_color, alpha=0.5)

    def calculate_percentage_deviation(self, values, instance_name):
        """Calculate percentage deviation from optimal value for best_candidates plots"""
        if self.config['plot_type'] != 'best_candidates' or instance_name not in self.objectives:
            return values

        optimal_value = self.objectives[instance_name]

        # Check if any values are significantly better than the known optimum
        # Allow small tolerance for numerical precision
        tolerance = 1e-4 * np.abs(optimal_value)
        significantly_better = values < (optimal_value - tolerance)
        # print(values - optimal_value)

        if np.any(significantly_better):
            better_indices = np.where(significantly_better)[0]
            best_found = np.min(values[better_indices])
            improvement = optimal_value - best_found
            improvement_pct = improvement / np.abs(optimal_value) * 100

            print(f"WARNING: Found {np.sum(significantly_better)} values better than known optimum for {instance_name}!")
            print(f"Known optimum: {optimal_value}")
            print(f"Best found: {best_found}")
            print(f"Improvement: {improvement} ({improvement_pct:.3f}%)")
            print(f"Using best found value as new reference point.")

            # Update the optimal value to the best found for this calculation
            optimal_value = best_found


        new_optimal_value = optimal_value
        if optimal_value < 0:
            new_optimal_value = - optimal_value
            values = values + 2*new_optimal_value

        # Calculate percentage deviation: (current - optimal) / |optimal| * 100
        # This will always be >= 0 since current >= optimal
        percentage_deviations = (values - new_optimal_value) / new_optimal_value * 100
        return percentage_deviations

    def load_data_for_algorithm_instance(self, algorithm, instance):
        csv_pattern = self.get_csv_pattern(self.config['plot_type'])

        algorithm_path = self.data_path / algorithm / instance
        csv_files = list(algorithm_path.glob(csv_pattern))

        if not csv_files:
            print(f"Warning: No CSV files found for {algorithm}/{instance}")
            return None

        print(f"Found {len(csv_files)} files for {algorithm}/{instance}")

        value_col = self.get_value_column(self.config['plot_type'])
        all_data = []

        for csv_file in csv_files:
            try:
                df = pd.read_csv(csv_file)
                if 'iteration' in df.columns and value_col in df.columns:
                    all_data.append(df[['iteration', value_col]].copy())
            except Exception as e:
                print(f"Error reading {csv_file}: {e}")

        if not all_data:
            return None

        # Find the maximum function evaluation (iteration value) across all runs
        max_evaluation = max(df['iteration'].max() for df in all_data)

        # Create a common evaluation grid - collect all unique evaluation points
        all_evaluations = set()
        for df in all_data:
            all_evaluations.update(df['iteration'].tolist())
        all_evaluations.add(max_evaluation)  # Ensure max is included
        common_grid = sorted(list(all_evaluations))

        aligned_data = []

        for df in all_data:
            df_copy = df.copy().sort_values('iteration').reset_index(drop=True)

            # Interpolate/extend this run to the common grid
            aligned_values = []
            current_value = df_copy[value_col].iloc[0] if len(df_copy) > 0 else 0
            df_idx = 0

            for eval_point in common_grid:
                # Update current_value if we have data at this evaluation point or before
                while df_idx < len(df_copy) and df_copy['iteration'].iloc[df_idx] <= eval_point:
                    current_value = df_copy[value_col].iloc[df_idx]
                    df_idx += 1

                aligned_values.append(current_value)

            aligned_data.append(aligned_values)

        # All runs now have the same length and evaluation points
        values_matrix = np.column_stack(aligned_data)
        values_matrix = values_matrix.astype(np.float64)
        iterations = np.array(common_grid)

        # Apply percentage deviation calculation for best_candidates
        if self.config['plot_type'] == 'best_candidates':
            # Apply percentage deviation to each column (run) separately
            for i in range(values_matrix.shape[1]):
                values_matrix[:, i] = self.calculate_percentage_deviation(values_matrix[:, i], instance)

        mean_values = np.mean(values_matrix, axis=1)
        std_values = np.std(values_matrix, axis=1)

        return {
            'iterations': iterations,
            'mean': mean_values,
            'std': std_values,
            'runs': len(all_data)
        }

    def create_plot(self):
        fig, ax = plt.subplots(1, 1, figsize=self.config['plot_settings']['figsize'])

        combination_method = self.config['plot_settings'].get('color_combination', 'blend')

        for instance in self.config['instances']:
            instance_name = instance['name'] if isinstance(instance, dict) else instance
            instance_label = instance.get('label', instance_name) if isinstance(instance, dict) else instance_name
            instance_color = instance.get('color', '#000000') if isinstance(instance, dict) else '#000000'

            for algorithm in self.config['algorithms']:
                alg_name = algorithm['name']
                alg_label = algorithm['label']
                alg_color = algorithm.get('color', '#000000')
                linestyle = algorithm['linestyle']

                # Get combined color
                combined_color = self.get_combined_color(alg_color, instance_color, combination_method)

                data = self.load_data_for_algorithm_instance(alg_name, instance_name)

                if data is None:
                    print(f"No data found for {alg_name}/{instance_name}")
                    continue

                # Handle data values (percentage deviation for best_candidates, absolute for others)
                mean_values = data['mean']
                std_values = data['std']

                # For log scale with negative values (non-best_candidates plots), take absolute value
                if self.config['plot_settings'].get('log_y', False) and self.config['plot_type'] != 'best_candidates':
                    mean_values = np.abs(mean_values)

                ax.plot(data['iterations'], mean_values,
                       color=combined_color,
                       linestyle=linestyle,
                       label=f"{alg_label} - {instance_label}",
                       linewidth=2,
                       drawstyle='steps-post')

                # Add fill_between for standard deviation bands (if enabled)
                if self.config['plot_settings'].get('show_std', True):
                    lower_bound = mean_values - std_values
                    upper_bound = mean_values + std_values

                    # For log scale, ensure positive values
                    if self.config['plot_settings'].get('log_y', False):
                        lower_bound = np.maximum(lower_bound, 0.01)

                    ax.fill_between(data['iterations'],
                                   lower_bound,
                                   upper_bound,
                                   color=combined_color,
                                   alpha=self.config['plot_settings']['alpha_fill'])

        ax.set_xlabel(self.config['plot_settings']['xlabel'], fontsize=14)
        ax.set_ylabel(self.config['plot_settings']['ylabel'], fontsize=14)
        ax.set_title(self.config['plot_settings']['title'], fontsize=16)

        # Set log scale if requested
        if self.config['plot_settings'].get('log_x', False):
            ax.set_xscale('log')
        if self.config['plot_settings'].get('log_y', False):
            ax.set_yscale('log')

        # Format y-axis with percentages for best_candidates plots
        if self.config['plot_type'] == 'best_candidates':
            ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x:.1f}%'))

        # Set tick label font sizes
        ax.tick_params(axis='both', which='major', labelsize=12)
        ax.tick_params(axis='both', which='minor', labelsize=10)

        if self.config['plot_settings']['grid']:
            ax.grid(True, alpha=0.3)

        if self.config['plot_settings']['legend']:
            ax.legend()

        plt.tight_layout()

        instance_names = [inst['name'] if isinstance(inst, dict) else inst for inst in self.config['instances']]
        output_file = self.output_dir / f"{self.config['plot_type']}_{'_'.join(instance_names)}.svg"
        plt.savefig(output_file, format='svg', bbox_inches='tight')
        print(f"Plot saved to: {output_file}")

def main():
    parser = argparse.ArgumentParser(description='Plot constraint optimization results')
    parser.add_argument('config', help='Path to JSON configuration file')
    args = parser.parse_args()

    plotter = ConstraintOptimizationPlotter(args.config)
    plotter.create_plot()

if __name__ == '__main__':
    main()