from matplotlib.backends.backend_qtagg import NavigationToolbar2QT
from PySide6.QtWidgets import QWidget, QVBoxLayout, QPushButton
from matplotlib.figure import Figure
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
import numpy as np
from acycle.util import utiltool
from acycle.util.utiltool import data_distribution
import pandas as pd
import os
from matplotlib.ticker import AutoMinorLocator
from scipy.stats import gaussian_kde
import matplotlib.font_manager as fm
@utiltool.window_shortcut
class MplWidget(QWidget):
    def __init__(self, parent=None, width=5, height=4, dpi=100, **kwargs):
        super().__init__(parent)
        self.file = kwargs.get('file_path', "")  # Default file list
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.canvas = FigureCanvasQTAgg(self.fig)
        self.toolbar = NavigationToolbar2QT(self.canvas, self)
        self.button = self.toolbar.addWidget(QPushButton("Button"))
        layout = QVBoxLayout(self)
        layout.addWidget(self.toolbar)
        layout.addWidget(self.canvas)
        
        # Initialize subplot system
        self.axes = {}  # Dictionary to store all subplots
        self.current_ax_key = '1'  # Key of the currently active subplot
        
        # Create default subplot
        self.ax = self.fig.add_subplot(111)
        self.axes['1'] = self.ax
        self.current_ax = self.ax
        
        # Default plot configuration
        self.plot_config = {
            'data': {
                'all_panels': '1,1',
                'set_panels': '1',
                'data_file': '',
                'x_column': 1,
                'y_column': 2,
                'set_data': ''
            },
            'title': {
                'title': '1D line plot',
                'legend': 'data1',
                'font_family': 'Arial',
                'font_weight': 'normal'
            },
            'axis': {
            'x_label': 'Depth (m)',
            'y_label': 'Value',  # Add y_label configuration
            'x_min': 0,
            'x_max': 2000,
            'y_min': 0,  # Add y_min configuration
            'y_max': 2000,  # Add y_max configuration
            'x_tick_label': False,
            'x_minor_tick': False,
            'y_minor_tick': True,
            'x_direction': 'normal',
            'y_direction': 'normal',  # Add y_direction configuration
            'x_log_scale': False,  # Changed to x_log_scale
            'y_log_scale': False   # Add y_log_scale
        },
            'legend': {
                'font_size': 12,
                'font_style': 'normal'
            },
            'line': {
                'plot_type': 'line',
                'line_style': '-',
                'line_width': 2,
                'line_color': '#1f77b4'
            },
            'marker': {
                'show_face': False,
                'show_edge': False,
                'style': 'o',
                'size': 1,
                'alpha': 0.5,
                'face_color': '#1f77b4',
                'edge_color': '#000000'
            }
        }

    def update_plot_config(self, config):
        """Update plot configuration and redraw"""
        self.plot_config.update(config)
        
        # Check if new subplot layout needs to be created
        panels_text = self.plot_config['data'].get('all_panels', '1,2')
        try:
            rows, cols = map(int, panels_text.split(','))
            total_panels = rows * cols
            
            # If number of subplots changed, recreate subplots
            if len(self.axes) != total_panels:
                self.create_subplots(rows, cols)
                
            # Set currently active subplot
            set_panel = self.plot_config['data'].get('set_panels', '1')
            if set_panel in self.axes:
                self.current_ax = self.axes[set_panel]
                self.current_ax_key = set_panel
            else:
                # If specified subplot doesn't exist, use first subplot
                self.current_ax = next(iter(self.axes.values()))
                self.current_ax_key = next(iter(self.axes.keys()))
                
        except ValueError:
            # If parsing fails, use default subplot
            if not self.axes:
                self.current_ax = self.fig.add_subplot(111)
                self.axes['1'] = self.current_ax
                self.current_ax_key = '1'
        
        self.redraw_plot()
    
    def create_subplots(self, rows, cols):
        """Create subplots with specified rows and columns"""
        # Clear all existing subplots
        self.fig.clear()
        self.axes.clear()
        
        # Create new subplots
        for i in range(1, rows * cols + 1):
            ax = self.fig.add_subplot(rows, cols, i)
            self.axes[str(i)] = ax
            
        # Set current active subplot to first one
        self.current_ax = next(iter(self.axes.values()))
        self.current_ax_key = next(iter(self.axes.keys()))
        
        # Redraw canvas
        self.canvas.draw()

    def redraw_plot(self):
        """Redraw graph according to current configuration"""
        if not self.file or not self.current_ax:
            return
        # Clear current active subplot
        self.current_ax.clear()
        
        # Select plot type based on configuration
        plot_type = self.plot_config['line'].get('plot_type', 'line')
        
        # Read data
        data = utiltool.load_data(self.file)
        x = data[:,0]
        y = data[:,1]
        # Draw graph
        if plot_type == 'bar':
            self.bar_plot(x, y, **self.get_style_kwargs())
        elif plot_type == 'scatter':
            self.scatter_plot(x, y, **self.get_style_kwargs())
        elif plot_type == 'distribution':
            self.plot_distribution()
        elif plot_type == 'sample_rate':
            self.plot_sample_rate()
        else:  # Default line plot
            self.plot(x, y, **self.get_style_kwargs())

        # Set title
        if self.plot_config['title'].get('title'):
            self.current_ax.set_title(
                self.plot_config['title']['title'],
                fontfamily=self.plot_config['title'].get('font_family', 'Arial'),
                fontweight=self.plot_config['title'].get('font_weight', 'normal')
            )
        
        # Set axis labels
        if self.plot_config['axis'].get('x_label'):
            self.current_ax.set_xlabel(self.plot_config['axis']['x_label'])
        
        if self.plot_config['axis'].get('y_label'):
            self.current_ax.set_ylabel(self.plot_config['axis']['y_label'])
        
        # Set axis range
        if (self.plot_config['axis'].get('x_min') is not None and 
            self.plot_config['axis'].get('x_max') is not None):
            self.current_ax.set_xlim(
                self.plot_config['axis']['x_min'], 
                self.plot_config['axis']['x_max']
            )
        
        if (self.plot_config['axis'].get('y_min') is not None and 
            self.plot_config['axis'].get('y_max') is not None):
            self.current_ax.set_ylim(
                self.plot_config['axis']['y_min'], 
                self.plot_config['axis']['y_max']
            )
        
        # Set axis direction
        if self.plot_config['axis'].get('x_direction') == 'reverse':
            self.current_ax.invert_xaxis()
        
        if self.plot_config['axis'].get('y_direction') == 'reverse':
            self.current_ax.invert_yaxis()
        
        # Set log scale
        if self.plot_config['axis'].get('x_log_scale'):
            self.current_ax.set_xscale('log')
        
        if self.plot_config['axis'].get('y_log_scale'):
            self.current_ax.set_yscale('log')
        
        # Set minor ticks
        if self.plot_config['axis'].get('x_minor_tick'):
            self.current_ax.xaxis.set_minor_locator(AutoMinorLocator())
        
        if self.plot_config['axis'].get('y_minor_tick'):
            self.current_ax.yaxis.set_minor_locator(AutoMinorLocator())
        
        # Set grid
        if self.plot_config['axis'].get('grid', True):
            self.current_ax.grid(True, linestyle='--', alpha=0.6)
        else:
            self.current_ax.grid(False)
        
        # Set legend
        if (self.plot_config['title'].get('legend') and 
            self.plot_config['legend'].get('show', True)):
            font_properties = fm.FontProperties(
                family=self.plot_config['title'].get('font_family', 'Arial'),
                weight=self.plot_config['title'].get('font_weight', 'normal'),
                size=self.plot_config['legend'].get('font_size', 12),
                style=self.plot_config['legend'].get('font_style', 'normal')
            )
            self.current_ax.legend([self.plot_config['title']['legend']], prop=font_properties)
        
        self.canvas.draw()

    def get_style_kwargs(self):
        """Extract style parameters from configuration"""
        style_kwargs = {}
        
        # Line style
        if 'line' in self.plot_config:
            line_config = self.plot_config['line']
            if 'line_color' in line_config:  # Fixed to line_color
                style_kwargs['color'] = line_config['line_color']
            if 'line_width' in line_config:
                style_kwargs['linewidth'] = line_config['line_width']
            if 'line_style' in line_config:
                style_kwargs['linestyle'] = line_config['line_style']
        
        # Marker style
        if 'marker' in self.plot_config:
            marker_config = self.plot_config['marker']
            if 'style' in marker_config and marker_config['style'] != 'none':
                style_kwargs['marker'] = marker_config['style']
            if 'size' in marker_config:
                style_kwargs['markersize'] = marker_config['size']
            if 'alpha' in marker_config:
                style_kwargs['alpha'] = marker_config['alpha']
            if 'face_color' in marker_config and marker_config.get('show_face', True):
                style_kwargs['markerfacecolor'] = marker_config['face_color']
            if 'edge_color' in marker_config and marker_config.get('show_edge', True):
                style_kwargs['markeredgecolor'] = marker_config['edge_color']
        
        return style_kwargs

    def bar_plot(self, x, y, **kwargs):
        """Draw bar chart"""
        if not self.current_ax:
            return
        
        # Convert plot-style kwargs to bar-style kwargs
        bar_kwargs = {}
        
        # Handle color parameter
        if 'color' in kwargs:
            bar_kwargs['color'] = kwargs['color']
        
        # Handle line width parameter
        if 'linewidth' in kwargs:
            bar_kwargs['linewidth'] = kwargs['linewidth']
        
        # Handle edge color parameter
        if 'edgecolor' in kwargs or 'markeredgecolor' in kwargs:
            edge_color = kwargs.get('edgecolor', kwargs.get('markeredgecolor', 'black'))
            bar_kwargs['edgecolor'] = edge_color
        
        # Handle alpha parameter
        if 'alpha' in kwargs:
            bar_kwargs['alpha'] = kwargs['alpha']
        
        # Handle line style parameter
        if 'linestyle' in kwargs:
            bar_kwargs['linestyle'] = kwargs['linestyle']
        
        # Set default width if not provided
        bar_kwargs['width'] = kwargs.get('width', 0.8)
        
        # Draw bar chart with converted parameters
        self.current_ax.bar(x, y, **bar_kwargs)
    
    def scatter_plot(self, x, y, **kwargs):
        """Draw scatter plot"""
        if not self.current_ax:
            return
        # Convert plot-style kwargs to scatter-style kwargs
        scatter_kwargs = {}
        # Handle color parameter
        if 'color' in kwargs:
            scatter_kwargs['c'] = kwargs['color']
        # Handle marker parameters
        if 'marker' in kwargs:
            scatter_kwargs['marker'] = kwargs['marker']
        # Handle size parameter (convert markersize to s for scatter)
        if 'markersize' in kwargs:
            scatter_kwargs['s'] = kwargs['markersize'] * 20  # Scale for better visibility
        # Handle alpha parameter
        if 'alpha' in kwargs:
            scatter_kwargs['alpha'] = kwargs['alpha']
        # Handle face and edge colors
        if 'markerfacecolor' in kwargs:
            scatter_kwargs['facecolors'] = kwargs['markerfacecolor']
        if 'markeredgecolor' in kwargs:
            scatter_kwargs['edgecolors'] = kwargs['markeredgecolor']
        # Draw scatter plot with converted parameters
        self.current_ax.scatter(x, y, **scatter_kwargs)
    
    def plot(self, x, y, **kwargs):
        if not self.current_ax:
            return
        self.current_ax.plot(x, y, **kwargs)
    
    def set_file_path(self, file):
        self.setWindowTitle(file)
        self.file = file
        
    def plot_sample_rate(self):
        """Draw depth-logFe relationship graph and annotate sampling characteristics"""
        # Main data curve
        # Load data
        data, avg_interval,_ = self.load_and_preprocess()
        # Step line drawing (staircase-like)
        self.ax.step(data['depth'], data['sample_interval'], 
            where='mid',
            color='b',
            linewidth=1.5,
            label='logFe (step)')
        # Annotate sampling points
        # ax.plot(data['logFe'], data['depth'], 'ro', markersize=4, alpha=0.5, label='Samples')
        # Set axes
        self.ax.set_xlabel('logFe Value', fontsize=12)
        self.ax.set_ylabel('Depth (ft)', fontsize=12)
        self.ax.set_title('Svalbard PETM - logFe Profile\nSampling Characteristics', fontsize=14, pad=20)
        self.ax.grid(True, linestyle='--', alpha=0.6)
        # Reverse depth axis (common in geological representation)
        # ax.invert_yaxis()
        # Add sampling rate information
        textstr = '\n'.join((
            f'Total Samples: {len(data)}',
            f'Depth Range: {data["depth"].min():.1f} - {data["depth"].max():.1f} ft',
            f'Avg. Sampling Interval: {avg_interval:.2f} ft',
            f'Avg. Sampling Rate: {1/avg_interval:.2f} samples/ft'
        ))
        props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
        self.ax.text(0.05, 0.95, textstr, transform=self.ax.transAxes, 
                verticalalignment='center_baseline', bbox=props, fontsize=10)
        # Add legend
        self.ax.legend(loc='lower right')
        # Beautify axes
        self.ax.xaxis.set_minor_locator(AutoMinorLocator())
        self.ax.yaxis.set_minor_locator(AutoMinorLocator())
    
    def load_and_preprocess(self):
        """Load data and calculate sampling interval"""
        data = utiltool.load_data(self.file)
        
        # Calculate sampling interval (difference between adjacent depths)
        depth_diff = np.float32(np.diff(data['depth']))

        avg_sample_interval = np.mean(depth_diff)
        # Fix: Add a NaN at the beginning to match length
        data['sample_interval'] = np.insert(depth_diff, 0, np.nan)
        return data, avg_sample_interval,depth_diff
    
    def plot_distribution(self):
        """Draw data distribution histogram"""
        # Load data
        _,_,data= self.load_and_preprocess()
        # Calculate statistics
        stats = {
            'max': np.max(data),
            'min': np.min(data),
            'mean': np.mean(data),
            'median': np.median(data),
            'variance': np.var(data)
        }
        # Create figure or use passed ax
        NUM = 20
        bin_width = (np.max(data) - np.min(data))/NUM
        bins = np.arange(np.min(data), np.max(data) + bin_width, bin_width)
        
        # Draw histogram
        n, bins, patches = self.ax.hist(data, bins, 
                                width=bin_width,
                                color='skyblue',
                                edgecolor='white',
                                alpha=0.7,
                                label=get_label(stats))
        
        # --- Draw KDE curve ---
        ax2 = self.ax.twinx()
        kde = gaussian_kde(data)
        x_grid = np.linspace(np.min(data)-1, np.max(data)+1, 500)
        y_kde = kde(x_grid)
        ax2.plot(x_grid, y_kde, 'r-', linewidth=2)
        
        # --- Graph beautification ---
        self.ax.set_title(os.path.basename(self.file)+': kernel fit of the data', pad=20)
        self.ax.set_xlabel('Value')
        self.ax.set_ylabel('Frequency (Histogram)', color='skyblue')
        ax2.set_ylabel('Density (KDE)', color='red')
        
        self.ax.tick_params(axis='y', colors='skyblue')
        ax2.tick_params(axis='y', colors='red')
        self.ax.grid(axis='y', alpha=0.3)
        
        # Set axis range
        self.ax.set_xlim(min(data)-bin_width, max(data)+bin_width)
        
        # Combine legends
        lines1, labels1 = self.ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        self.ax.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        # Display control
        self.fig.tight_layout()
        if hasattr(self.fig, 'canvas'):
            self.fig.canvas.draw()
    
def get_label(stats):
    return  f"Max: {stats['max']}, Min: {stats['min']}, Mean: {stats['mean']}, Median: {stats['median']}, Variance: {stats['variance']}"