#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Dispersion Curve Inversion Visualization Utilities

Plotting functions for dispersion curves and inversion results.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
import warnings


class DCInversionVisualizations:
    """Container for visualization methods"""
    
    @staticmethod
    def plot_dispersion_data(ax, positions_data, selected_only=True,
                            mode_filter='All', color_by='mode', show_errors=True, 
                            x_axis='frequency', cmap='viridis'):
        """
        Plot dispersion curves for multiple modes with flexible options
        
        Parameters:
        -----------
        ax : matplotlib axis
        positions_data : dict
            Position data dictionary
        selected_only : bool
            If True, only plot positions marked for inversion
        mode_filter : str
            'All' or specific mode name ('M0', 'M1', etc)
        color_by : str
            'mode' or 'position' - how to color the curves
        show_errors : bool
            If True, show error bars
        x_axis : str
            'frequency' or 'wavelength' - x-axis variable
        cmap : str
            Colormap name for position coloring
        """
        from matplotlib import cm
        from matplotlib.colors import Normalize
        from matplotlib.colorbar import ColorbarBase
        
        # Define mode colors
        mode_colors = {
            'M0': '#1f77b4',  # blue
            'M1': '#ff7f0e',  # orange
            'M2': '#2ca02c',  # green
            'M3': '#d62728',  # red
            'M4': '#9467bd',  # purple
            'M5': '#8c564b',  # brown
            'M6': '#e377c2',  # pink
        }
        
        # Collect all positions and modes for colormap
        all_xmids = sorted(positions_data.keys())
        all_modes = set()
        for pos_data in positions_data.values():
            all_modes.update(pos_data.get('modes', {}).keys())
        all_modes = sorted(all_modes, key=lambda x: int(x.replace('M', '')))
        
        # Create position colormap
        cmap_obj = None
        norm = None
        if color_by == 'position' and len(all_xmids) > 1:
            cmap_obj = cm.get_cmap(cmap)
            norm = Normalize(vmin=min(all_xmids), vmax=max(all_xmids))
            pos_to_color = {xmid: cmap_obj(norm(xmid)) for xmid in all_xmids}
        
        # Set x-axis label based on parameter
        x_label = 'Wavelength (m)' if x_axis == 'wavelength' else 'Frequency (Hz)'
        
        # Plot data
        legend_added = set()  # Track which legend entries we've added
        
        for xmid in all_xmids:
            pos_data = positions_data[xmid]
            
            # Filter by position selection
            if selected_only and not pos_data.get('selected_for_inversion', True):
                continue
            
            for mode_name in all_modes:
                if mode_name not in pos_data.get('modes', {}):
                    continue
                
                # Apply mode filter
                if mode_filter != "All" and mode_name != mode_filter:
                    continue
                
                mode_data = pos_data['modes'][mode_name]
                
                # Use resampled data if available
                data = mode_data.get('resampled', mode_data.get('original'))
                if data is None:
                    continue
                
                freq = data.get('frequencies', [])
                vel = data.get('velocities', [])
                err = data.get('errors', None)
                
                if len(freq) == 0:
                    continue
                
                # Convert to numpy arrays
                freq = np.asarray(freq)
                vel = np.asarray(vel)
                
                # Handle errors - robust check for None, empty, or invalid data
                try:
                    if err is None:
                        err = np.array([])
                    else:
                        err = np.asarray(err)
                        # Check if empty after conversion
                        if err.size == 0:
                            err = np.array([])
                        else:
                            # Ensure error values are non-negative (required by matplotlib)
                            err = np.abs(err)
                except (ValueError, TypeError, AttributeError):
                    # If anything fails, use empty array
                    err = np.array([])
                
                # Calculate x-axis data
                if x_axis == 'wavelength':
                    x_data = vel / freq  # wavelength = velocity / frequency
                else:
                    x_data = freq
                
                # Determine color and label
                if color_by == 'position':
                    color = pos_to_color[xmid]
                    # No label when using colorbar
                    label = None
                else:  # color by mode
                    color = mode_colors.get(mode_name, 'gray')
                    label_key = mode_name
                    # Only add legend entry once per mode
                    if label_key not in legend_added:
                        legend_added.add(label_key)
                        label = label_key
                    else:
                        label = None
                
                # Plot with or without error bars
                if show_errors and len(err) > 0 and np.any(err > 0):
                    ax.errorbar(x_data, vel, yerr=err, fmt='o', color=color,
                               label=label, markersize=5, alpha=0.8, capsize=3,
                               elinewidth=1.5, capthick=1.5)
                else:
                    ax.plot(x_data, vel, 'o', color=color, label=label,
                           markersize=5, alpha=0.8)
        
        ax.set_xlabel(x_label)
        ax.set_ylabel('Phase Velocity (m/s)')
        ax.grid(True, alpha=0.3)
        
        # Add legend for mode coloring, colorbar for position coloring
        if color_by == 'position' and len(all_xmids) > 1 and cmap_obj is not None:
            # Add colorbar instead of legend
            sm = cm.ScalarMappable(cmap=cmap_obj, norm=norm)
            sm.set_array([])
            cbar = ax.figure.colorbar(sm, ax=ax)
            cbar.set_label('Position (m)', rotation=270, labelpad=15)
        elif len(legend_added) > 0:
            # Add legend for mode coloring
            ax.legend(fontsize=8, loc='best', ncol=min(3, len(legend_added)))
    
    @staticmethod
    def plot_vs_profile(ax, vs_model, thickness, label='Inverted', color='blue',
                        show_layers=True):
        """
        Plot 1D Vs profile
        
        Parameters:
        -----------
        ax : matplotlib axis
        vs_model : array
            Vs values for each layer
        thickness : array
            Layer thicknesses
        label : str
            Label for legend
        color : str
            Line color
        show_layers : bool
            If True, show layer boundaries
        """
        # Calculate depths (top of each layer)
        depths = np.concatenate([[0], np.cumsum(thickness)])
        max_depth = depths[-1] + thickness[-1] * 0.5  # Extend half-space
        depths = np.append(depths, max_depth)
        
        # Extend vs_model for half-space
        vs_extended = np.append(vs_model, vs_model[-1])
        
        # Create step plot
        ax.step(vs_extended, depths, where='post', label=label, 
                color=color, linewidth=2)
        
        # Show layer boundaries
        if show_layers:
            for d in depths[1:-1]:
                ax.axhline(y=d, color='gray', linestyle='--', 
                          alpha=0.3, linewidth=0.5)
        
        ax.set_xlabel('Vs (m/s)')
        ax.set_ylabel('Depth (m)')
        ax.invert_yaxis()
        ax.grid(True, alpha=0.3)
        ax.legend()
    
    @staticmethod
    def plot_dispersion_fit(ax, mode_fits, mode_name):
        """
        Plot observed vs calculated dispersion for one mode
        
        Parameters:
        -----------
        ax : matplotlib axis
        mode_fits : dict
            Mode fit results from inversion
        mode_name : str
            Which mode to plot
        """
        if mode_name not in mode_fits:
            ax.text(0.5, 0.5, f'No data for {mode_name}',
                   ha='center', va='center', transform=ax.transAxes)
            return
        
        fit_data = mode_fits[mode_name]
        
        # Get frequency axis (need to reconstruct from position_data if available)
        # For now, use index
        n = len(fit_data['observed'])
        x_axis = np.arange(n)
        
        # Plot observed and calculated
        ax.plot(x_axis, fit_data['observed'], 'ko', label='Observed',
               markersize=6, alpha=0.7)
        ax.plot(x_axis, fit_data['calculated'], 'r-', label='Calculated',
               linewidth=2)
        
        # Add RMS and chi2 to title
        rms = fit_data['rms']
        chi2 = fit_data['chi2']
        ax.set_title(f'{mode_name}: RMS={rms:.2f} m/s, χ²={chi2:.2f}')
        
        ax.set_xlabel('Data Point Index')
        ax.set_ylabel('Phase Velocity (m/s)')
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    @staticmethod
    def plot_residuals(ax, mode_fits, mode_name):
        """
        Plot residuals for one mode
        
        Parameters:
        -----------
        ax : matplotlib axis
        mode_fits : dict
            Mode fit results
        mode_name : str
            Which mode to plot
        """
        if mode_name not in mode_fits:
            return
        
        fit_data = mode_fits[mode_name]
        residuals = fit_data['residuals']
        n = len(residuals)
        x_axis = np.arange(n)
        
        ax.bar(x_axis, residuals, color='gray', alpha=0.7)
        ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
        
        ax.set_xlabel('Data Point Index')
        ax.set_ylabel('Residual (m/s)')
        ax.set_title(f'{mode_name} Residuals')
        ax.grid(True, alpha=0.3)
    
    @staticmethod
    def plot_all_modes_fit(fig, mode_fits, position_data=None):
        """
        Create multi-panel plot for all modes
        
        Parameters:
        -----------
        fig : matplotlib Figure
        mode_fits : dict
            All mode fit results
        position_data : dict, optional
            Position data with frequency information
        """
        n_modes = len(mode_fits)
        if n_modes == 0:
            return
        
        # Create subplots: 2 columns (fit and residuals)
        fig.clear()
        axes = fig.subplots(n_modes, 2, figsize=(12, 3 * n_modes))
        
        if n_modes == 1:
            axes = axes.reshape(1, 2)
        
        for idx, (mode_name, fit_data) in enumerate(sorted(mode_fits.items())):
            # Left: observed vs calculated
            ax_fit = axes[idx, 0]
            
            # Try to get actual frequencies if position_data provided
            if position_data and 'modes' in position_data and mode_name in position_data['modes']:
                mode_data = position_data['modes'][mode_name]
                data = mode_data.get('resampled', mode_data.get('original'))
                if data and 'frequencies' in data:
                    x_axis = np.asarray(data['frequencies'])
                    xlabel = 'Frequency (Hz)'
                else:
                    x_axis = np.arange(len(fit_data['observed']))
                    xlabel = 'Data Point Index'
            else:
                x_axis = np.arange(len(fit_data['observed']))
                xlabel = 'Data Point Index'
            
            ax_fit.plot(x_axis, fit_data['observed'], 'ko', 
                       label='Observed', markersize=6, alpha=0.7)
            ax_fit.plot(x_axis, fit_data['calculated'], 'r-',
                       label='Calculated', linewidth=2)
            
            rms = fit_data['rms']
            chi2 = fit_data['chi2']
            ax_fit.set_title(f'{mode_name}: RMS={rms:.2f} m/s, χ²={chi2:.2f}')
            ax_fit.set_xlabel(xlabel)
            ax_fit.set_ylabel('Phase Velocity (m/s)')
            ax_fit.legend()
            ax_fit.grid(True, alpha=0.3)
            
            # Right: residuals
            ax_res = axes[idx, 1]
            ax_res.bar(x_axis, fit_data['residuals'], color='gray', alpha=0.7)
            ax_res.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
            ax_res.set_title(f'{mode_name} Residuals')
            ax_res.set_xlabel(xlabel)
            ax_res.set_ylabel('Residual (m/s)')
            ax_res.grid(True, alpha=0.3)
        
        fig.tight_layout()
    
    @staticmethod
    def plot_batch_results(fig, all_results):
        """
        Plot comparison of results from multiple positions
        
        Parameters:
        -----------
        fig : matplotlib Figure
        all_results : dict
            {xmid: result_dict, ...}
        """
        fig.clear()
        
        # Filter successful results
        successful = {xmid: res for xmid, res in all_results.items() 
                     if res.get('success', False)}
        
        if len(successful) == 0:
            ax = fig.add_subplot(111)
            ax.text(0.5, 0.5, 'No successful inversions to display',
                   ha='center', va='center', transform=ax.transAxes)
            return
        
        # Create 3 subplots
        axes = fig.subplots(1, 3)
        
        # Plot 1: All Vs profiles overlaid
        ax1 = axes[0]
        colors = plt.cm.viridis(np.linspace(0, 1, len(successful)))
        
        for (xmid, result), color in zip(sorted(successful.items()), colors):
            vs_model = result['vs_model']
            thickness = result['thickness']
            
            depths = np.concatenate([[0], np.cumsum(thickness)])
            max_depth = depths[-1] + thickness[-1] * 0.5
            depths = np.append(depths, max_depth)
            vs_extended = np.append(vs_model, vs_model[-1])
            
            ax1.step(vs_extended, depths, where='post', 
                    label=f'{xmid:.1f}m', color=color, linewidth=2)
        
        ax1.set_xlabel('Vs (m/s)')
        ax1.set_ylabel('Depth (m)')
        ax1.invert_yaxis()
        ax1.legend(fontsize=8)
        ax1.set_title('All Vs Profiles')
        ax1.grid(True, alpha=0.3)
        
        # Plot 2: Chi² comparison
        ax2 = axes[1]
        xmids = []
        chi2s = []
        colors_chi2 = []
        
        for xmid, result in sorted(successful.items()):
            xmids.append(xmid)
            chi2 = result['total_chi2']
            chi2s.append(chi2)
            # Color based on quality: green if chi2 < 1, yellow if < 2, red otherwise
            if chi2 < 1.0:
                colors_chi2.append('green')
            elif chi2 < 2.0:
                colors_chi2.append('orange')
            else:
                colors_chi2.append('red')
        
        bars = ax2.bar(range(len(xmids)), chi2s, color=colors_chi2, alpha=0.7)
        ax2.set_xticks(range(len(xmids)))
        ax2.set_xticklabels([f'{x:.1f}' for x in xmids], rotation=45)
        ax2.set_xlabel('Position (m)')
        ax2.set_ylabel('χ²')
        ax2.set_title('Inversion Quality')
        ax2.axhline(y=1.0, color='darkgreen', linestyle='--', 
                   linewidth=2, label='Target χ²=1')
        ax2.legend()
        ax2.grid(True, alpha=0.3, axis='y')
        
        # Plot 3: RMS comparison
        ax3 = axes[2]
        rms_values = [result['total_rms'] for result in 
                     [successful[x] for x in sorted(successful.keys())]]
        
        ax3.bar(range(len(xmids)), rms_values, color='steelblue', alpha=0.7)
        ax3.set_xticks(range(len(xmids)))
        ax3.set_xticklabels([f'{x:.1f}' for x in xmids], rotation=45)
        ax3.set_xlabel('Position (m)')
        ax3.set_ylabel('RMS (m/s)')
        ax3.set_title('RMS Error')
        ax3.grid(True, alpha=0.3, axis='y')
        
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', UserWarning)
            fig.tight_layout()
    
    @staticmethod
    def create_summary_table_data(result):
        """
        Generate summary statistics table data
        
        Parameters:
        -----------
        result : dict
            Inversion result
        
        Returns:
        --------
        list of lists
            Table data [[row1], [row2], ...]
        """
        if not result.get('success', False):
            return [['Error', result.get('error', 'Unknown error')]]
        
        table_data = []
        table_data.append(['Overall Statistics', ''])
        table_data.append(['Total RMS', f"{result['total_rms']:.2f} m/s"])
        table_data.append(['Relative RMS', f"{result['rrms']:.2f} %"])
        table_data.append(['Chi-squared', f"{result['total_chi2']:.3f}"])
        table_data.append(['Iterations', str(result.get('iterations', 'N/A'))])
        table_data.append(['Converged', 'Yes' if result.get('convergence', False) else 'No'])
        table_data.append(['', ''])
        
        # Per-mode statistics
        table_data.append(['Mode Statistics', ''])
        for mode_name, fit_data in sorted(result.get('mode_fits', {}).items()):
            table_data.append([f'{mode_name} - RMS', f"{fit_data['rms']:.2f} m/s"])
            table_data.append([f'{mode_name} - Chi²', f"{fit_data['chi2']:.3f}"])
            table_data.append([f'{mode_name} - N points', str(fit_data['n'])])
        
        return table_data
