"""
Plotting Module for HydroAnomaly

This module provides easy-to-use plotting functions for USGS water data time series.
Creates professional-looking plots with minimal code.
"""

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import seaborn as sns
import pandas as pd
import numpy as np
from datetime import datetime
from typing import Optional, Tuple, List, Dict, Any
import warnings

# Set style
plt.style.use('default')
sns.set_palette("husl")


class WaterDataPlotter:
    """
    A class for creating professional time series plots of water data.
    
    This class provides methods to create various types of plots including
    basic time series, multi-parameter plots, and statistical visualizations.
    """
    
    def __init__(self, style: str = 'seaborn-v0_8', figsize: Tuple[int, int] = (12, 6)):
        """
        Initialize the plotter with default settings.
        
        Args:
            style (str): Matplotlib style to use
            figsize (tuple): Default figure size (width, height)
        """
        self.default_figsize = figsize
        self.colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                      '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        
        # Set plotting style
        try:
            plt.style.use(style)
        except:
            plt.style.use('default')
            warnings.warn(f"Style '{style}' not available, using default")
    
    def plot_timeseries(
        self,
        data: pd.DataFrame,
        parameter_name: str = "Value",
        title: Optional[str] = None,
        ylabel: Optional[str] = None,
        color: str = '#1f77b4',
        save_path: Optional[str] = None,
        show_stats: bool = True,
        figsize: Optional[Tuple[int, int]] = None
    ) -> plt.Figure:
        """
        Create a basic time series plot.
        
        Args:
            data (pd.DataFrame): Data with 'datetime' and 'value' columns
            parameter_name (str): Name of the parameter being plotted
            title (str, optional): Plot title
            ylabel (str, optional): Y-axis label
            color (str): Line color
            save_path (str, optional): Path to save the plot
            show_stats (bool): Whether to show statistics on the plot
            figsize (tuple, optional): Figure size
            
        Returns:
            matplotlib.figure.Figure: The created figure
        """
        if len(data) == 0:
            raise ValueError("No data to plot")
        
        figsize = figsize or self.default_figsize
        fig, ax = plt.subplots(figsize=figsize)
        
        # Plot the data
        ax.plot(data['datetime'], data['value'], color=color, linewidth=1.5, alpha=0.8)
        
        # Customize the plot
        if title is None:
            title = f"{parameter_name} Time Series"
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
        
        if ylabel is None:
            ylabel = parameter_name
        ax.set_ylabel(ylabel, fontsize=12)
        ax.set_xlabel('Date', fontsize=12)
        
        # Format dates on x-axis
        self._format_date_axis(ax, data['datetime'])
        
        # Add grid
        ax.grid(True, alpha=0.3, linestyle='--')
        
        # Add statistics if requested
        if show_stats:
            self._add_statistics_text(ax, data['value'], parameter_name)
        
        # Improve layout
        plt.tight_layout()
        
        # Save if requested
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"📊 Plot saved to: {save_path}")
        
        return fig
    
    def plot_multiple_parameters(
        self,
        data_dict: Dict[str, pd.DataFrame],
        title: str = "Multiple Parameters Time Series",
        save_path: Optional[str] = None,
        figsize: Optional[Tuple[int, int]] = None
    ) -> plt.Figure:
        """
        Plot multiple parameters on separate subplots.
        
        Args:
            data_dict (dict): Dictionary with parameter names as keys and DataFrames as values
            title (str): Main plot title
            save_path (str, optional): Path to save the plot
            figsize (tuple, optional): Figure size
            
        Returns:
            matplotlib.figure.Figure: The created figure
        """
        n_params = len(data_dict)
        if n_params == 0:
            raise ValueError("No data provided")
        
        figsize = figsize or (12, 4 * n_params)
        fig, axes = plt.subplots(n_params, 1, figsize=figsize, sharex=True)
        
        if n_params == 1:
            axes = [axes]
        
        colors = self.colors[:n_params]
        
        for i, (param_name, data) in enumerate(data_dict.items()):
            if len(data) == 0:
                continue
                
            ax = axes[i]
            ax.plot(data['datetime'], data['value'], 
                   color=colors[i], linewidth=1.5, alpha=0.8, label=param_name)
            
            ax.set_ylabel(param_name, fontsize=11)
            ax.grid(True, alpha=0.3, linestyle='--')
            ax.legend(loc='upper right')
            
            # Add basic stats
            mean_val = data['value'].mean()
            ax.axhline(y=mean_val, color=colors[i], linestyle=':', alpha=0.6, 
                      label=f'Mean: {mean_val:.2f}')
        
        # Format the bottom subplot x-axis
        if data_dict:
            sample_data = next(iter(data_dict.values()))
            self._format_date_axis(axes[-1], sample_data['datetime'])
        
        axes[-1].set_xlabel('Date', fontsize=12)
        fig.suptitle(title, fontsize=14, fontweight='bold')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"📊 Plot saved to: {save_path}")
        
        return fig
    
    def plot_comparison(
        self,
        data_list: List[Tuple[pd.DataFrame, str]],
        title: str = "Data Comparison",
        ylabel: str = "Value",
        save_path: Optional[str] = None,
        figsize: Optional[Tuple[int, int]] = None
    ) -> plt.Figure:
        """
        Plot multiple datasets on the same axes for comparison.
        
        Args:
            data_list (list): List of tuples (DataFrame, label)
            title (str): Plot title
            ylabel (str): Y-axis label
            save_path (str, optional): Path to save the plot
            figsize (tuple, optional): Figure size
            
        Returns:
            matplotlib.figure.Figure: The created figure
        """
        if not data_list:
            raise ValueError("No data provided")
        
        figsize = figsize or self.default_figsize
        fig, ax = plt.subplots(figsize=figsize)
        
        colors = self.colors[:len(data_list)]
        
        for i, (data, label) in enumerate(data_list):
            if len(data) == 0:
                continue
            
            ax.plot(data['datetime'], data['value'], 
                   color=colors[i], linewidth=1.5, alpha=0.8, label=label)
        
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
        ax.set_ylabel(ylabel, fontsize=12)
        ax.set_xlabel('Date', fontsize=12)
        ax.grid(True, alpha=0.3, linestyle='--')
        ax.legend()
        
        # Format dates
        if data_list and len(data_list[0][0]) > 0:
            self._format_date_axis(ax, data_list[0][0]['datetime'])
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"📊 Plot saved to: {save_path}")
        
        return fig
    
    def plot_statistics(
        self,
        data: pd.DataFrame,
        parameter_name: str = "Parameter",
        save_path: Optional[str] = None
    ) -> plt.Figure:
        """
        Create statistical plots (histogram and box plot).
        
        Args:
            data (pd.DataFrame): Data with 'datetime' and 'value' columns
            parameter_name (str): Name of the parameter
            save_path (str, optional): Path to save the plot
            
        Returns:
            matplotlib.figure.Figure: The created figure
        """
        if len(data) == 0:
            raise ValueError("No data to plot")
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        # Histogram
        ax1.hist(data['value'], bins=30, alpha=0.7, color=self.colors[0], edgecolor='black')
        ax1.set_title(f'{parameter_name} Distribution', fontweight='bold')
        ax1.set_xlabel(parameter_name)
        ax1.set_ylabel('Frequency')
        ax1.grid(True, alpha=0.3)
        
        # Add statistics to histogram
        mean_val = data['value'].mean()
        median_val = data['value'].median()
        ax1.axvline(mean_val, color='red', linestyle='--', alpha=0.8, label=f'Mean: {mean_val:.2f}')
        ax1.axvline(median_val, color='orange', linestyle='--', alpha=0.8, label=f'Median: {median_val:.2f}')
        ax1.legend()
        
        # Box plot
        box_data = ax2.boxplot(data['value'], patch_artist=True)
        box_data['boxes'][0].set_facecolor(self.colors[1])
        box_data['boxes'][0].set_alpha(0.7)
        
        ax2.set_title(f'{parameter_name} Box Plot', fontweight='bold')
        ax2.set_ylabel(parameter_name)
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"📊 Plot saved to: {save_path}")
        
        return fig
    
    def _format_date_axis(self, ax, dates):
        """Format the date axis based on the date range."""
        date_range = (dates.max() - dates.min()).days
        
        if date_range <= 7:  # Less than a week
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%m/%d %H:%M'))
            ax.xaxis.set_major_locator(mdates.HourLocator(interval=6))
        elif date_range <= 31:  # Less than a month
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%m/%d'))
            ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))
        elif date_range <= 365:  # Less than a year
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
            ax.xaxis.set_major_locator(mdates.MonthLocator())
        else:  # More than a year
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
            ax.xaxis.set_major_locator(mdates.YearLocator())
        
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')
    
    def _add_statistics_text(self, ax, values, parameter_name):
        """Add statistics text box to the plot."""
        stats_text = (
            f"Statistics:\n"
            f"Mean: {values.mean():.2f}\n"
            f"Median: {values.median():.2f}\n"
            f"Min: {values.min():.2f}\n"
            f"Max: {values.max():.2f}\n"
            f"Std: {values.std():.2f}"
        )
        
        # Position the text box
        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=9,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))


# Convenience functions for easy plotting
def plot_usgs_data(
    data: pd.DataFrame,
    parameter_name: str = "Value",
    title: Optional[str] = None,
    save_path: Optional[str] = None,
    show_stats: bool = True
) -> plt.Figure:
    """
    Quick function to plot USGS time series data.
    
    Args:
        data (pd.DataFrame): Data with 'datetime' and 'value' columns
        parameter_name (str): Name of the parameter being plotted
        title (str, optional): Plot title
        save_path (str, optional): Path to save the plot
        show_stats (bool): Whether to show statistics on the plot
        
    Returns:
        matplotlib.figure.Figure: The created figure
        
    Example:
        >>> import hydroanomaly
        >>> data = hydroanomaly.get_discharge("08158000", "2023-01-01", "2023-01-31")
        >>> hydroanomaly.plot_usgs_data(data, "Discharge (cfs)", "Colorado River Discharge")
    """
    plotter = WaterDataPlotter()
    return plotter.plot_timeseries(
        data=data,
        parameter_name=parameter_name,
        title=title,
        save_path=save_path,
        show_stats=show_stats
    )


def plot_multiple_gages(
    data_dict: Dict[str, pd.DataFrame],
    title: str = "Multiple Gage Comparison",
    parameter_name: str = "Value",
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Plot data from multiple gages for comparison.
    
    Args:
        data_dict (dict): Dictionary with gage IDs as keys and DataFrames as values
        title (str): Plot title
        parameter_name (str): Y-axis label
        save_path (str, optional): Path to save the plot
        
    Returns:
        matplotlib.figure.Figure: The created figure
    """
    plotter = WaterDataPlotter()
    data_list = [(data, f"Gage {gage_id}") for gage_id, data in data_dict.items()]
    return plotter.plot_comparison(data_list, title, parameter_name, save_path)


def quick_plot(data: pd.DataFrame, title: str = "USGS Data") -> None:
    """
    Create a quick plot and show it immediately.
    
    Args:
        data (pd.DataFrame): Data to plot
        title (str): Plot title
    """
    plot_usgs_data(data, title=title)
    plt.show()
