"""
Time series visualization for cave monitoring data.
Creates plots of V, L, F parameters and derived quantities over time.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

try:
    import matplotlib.pyplot as plt
    from matplotlib import dates as mdates
    from matplotlib.patches import Rectangle
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    logger.warning("Matplotlib not available - time series plots limited")

try:
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    PLOTLY_AVAILABLE = True
except ImportError:
    PLOTLY_AVAILABLE = False


class TimeSeriesPlots:
    """
    Creates time series visualizations for cave monitoring data.
    """
    
    def __init__(self):
        """Initialize time series plots."""
        self.plots = []
        logger.info("Initialized TimeSeriesPlots")

    def plot_vlf_series(self,
                       times: np.ndarray,
                       V: np.ndarray,
                       L: np.ndarray,
                       F: np.ndarray,
                       title: str = "VLF Time Series",
                       highlight_events: Optional[List[Tuple[float, float]]] = None) -> Optional[plt.Figure]:
        """
        Plot V, L, F parameters over time.
        
        Args:
            times: Time array
            V: Vertical passage values
            L: Lateral conduit values
            F: Flowstone values
            title: Plot title
            highlight_events: List of (start, end) times to highlight
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10), sharex=True)
        
        # Plot V
        ax1.plot(times, V, 'b-', linewidth=1.5, label='V')
        ax1.set_ylabel('Vertical Passage (m)')
        ax1.grid(True, alpha=0.3)
        ax1.legend(loc='upper right')
        
        # Plot L
        ax2.plot(times, L, 'r-', linewidth=1.5, label='L')
        ax2.set_ylabel('Lateral Conduit (m)')
        ax2.grid(True, alpha=0.3)
        ax2.legend(loc='upper right')
        
        # Plot F
        ax3.plot(times, F, 'g-', linewidth=1.5, label='F')
        ax3.set_ylabel('Flowstone (mm)')
        ax3.set_xlabel('Time')
        ax3.grid(True, alpha=0.3)
        ax3.legend(loc='upper right')
        
        # Highlight events
        if highlight_events:
            for start, end in highlight_events:
                for ax in [ax1, ax2, ax3]:
                    ax.axvspan(start, end, alpha=0.2, color='yellow')
        
        plt.suptitle(title)
        plt.tight_layout()
        
        self.plots.append(fig)
        return fig

    def plot_derivatives(self,
                        times: np.ndarray,
                        dVdt: np.ndarray,
                        dLdt: np.ndarray,
                        dFdt: np.ndarray,
                        title: str = "Rate of Change") -> Optional[plt.Figure]:
        """
        Plot time derivatives of V, L, F.
        
        Args:
            times: Time array
            dVdt, dLdt, dFdt: Derivatives
            title: Plot title
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        fig, axes = plt.subplots(3, 1, figsize=(12, 8), sharex=True)
        
        # Plot derivatives
        axes[0].plot(times, dVdt, 'b-', linewidth=1)
        axes[0].axhline(y=0, color='k', linestyle='--', linewidth=0.5)
        axes[0].set_ylabel('dV/dt (m/yr)')
        axes[0].grid(True, alpha=0.3)
        
        axes[1].plot(times, dLdt, 'r-', linewidth=1)
        axes[1].axhline(y=0, color='k', linestyle='--', linewidth=0.5)
        axes[1].set_ylabel('dL/dt (m/yr)')
        axes[1].grid(True, alpha=0.3)
        
        axes[2].plot(times, dFdt, 'g-', linewidth=1)
        axes[2].axhline(y=0, color='k', linestyle='--', linewidth=0.5)
        axes[2].set_ylabel('dF/dt (mm/yr)')
        axes[2].set_xlabel('Time')
        axes[2].grid(True, alpha=0.3)
        
        plt.suptitle(title)
        plt.tight_layout()
        
        self.plots.append(fig)
        return fig

    def plot_correlation_matrix(self,
                               data_dict: Dict[str, np.ndarray],
                               title: str = "Parameter Correlations") -> Optional[plt.Figure]:
        """
        Plot correlation matrix of parameters.
        
        Args:
            data_dict: Dictionary of parameter name to values
            title: Plot title
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        # Create data matrix
        names = list(data_dict.keys())
        n_vars = len(names)
        data = np.array([data_dict[name] for name in names])
        
        # Calculate correlation matrix
        corr_matrix = np.corrcoef(data)
        
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # Create heatmap
        im = ax.imshow(corr_matrix, cmap='RdBu_r', vmin=-1, vmax=1, aspect='auto')
        plt.colorbar(im, ax=ax, label='Correlation')
        
        # Add labels
        ax.set_xticks(np.arange(n_vars))
        ax.set_yticks(np.arange(n_vars))
        ax.set_xticklabels(names)
        ax.set_yticklabels(names)
        
        # Add correlation values
        for i in range(n_vars):
            for j in range(n_vars):
                text = ax.text(j, i, f'{corr_matrix[i, j]:.2f}',
                             ha="center", va="center", color="black" if abs(corr_matrix[i, j]) < 0.5 else "white")
        
        plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
        ax.set_title(title)
        
        plt.tight_layout()
        
        self.plots.append(fig)
        return fig

    def plot_seasonal_decomposition(self,
                                   times: np.ndarray,
                                   values: np.ndarray,
                                   parameter_name: str,
                                   period: float = 1.0) -> Optional[plt.Figure]:
        """
        Plot seasonal decomposition of time series.
        
        Args:
            times: Time array
            values: Parameter values
            parameter_name: Name of parameter
            period: Seasonal period (years)
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        from scipy import signal
        
        # Detrend
        coeffs = np.polyfit(times - times[0], values, 1)
        trend = np.polyval(coeffs, times - times[0])
        detrended = values - trend
        
        # Extract seasonal component
        seasonal = np.zeros_like(times)
        n_periods = int(np.ceil((times[-1] - times[0]) / period))
        
        for i in range(n_periods):
            start = times[0] + i * period
            end = start + period
            mask = (times >= start) & (times < end)
            if np.any(mask):
                seasonal[mask] = np.mean(detrended[mask])
        
        # Residual
        residual = detrended - seasonal
        
        fig, axes = plt.subplots(4, 1, figsize=(12, 10), sharex=True)
        
        # Original
        axes[0].plot(times, values, 'b-', linewidth=1)
        axes[0].set_ylabel('Original')
        axes[0].grid(True, alpha=0.3)
        
        # Trend
        axes[1].plot(times, trend, 'r-', linewidth=1)
        axes[1].set_ylabel('Trend')
        axes[1].grid(True, alpha=0.3)
        
        # Seasonal
        axes[2].plot(times, seasonal, 'g-', linewidth=1)
        axes[2].set_ylabel('Seasonal')
        axes[2].grid(True, alpha=0.3)
        
        # Residual
        axes[3].plot(times, residual, 'k-', linewidth=0.5)
        axes[3].axhline(y=0, color='gray', linestyle='--', linewidth=0.5)
        axes[3].set_ylabel('Residual')
        axes[3].set_xlabel('Time')
        axes[3].grid(True, alpha=0.3)
        
        plt.suptitle(f'{parameter_name} - Seasonal Decomposition')
        plt.tight_layout()
        
        self.plots.append(fig)
        return fig

    def plot_forecast(self,
                     historical_times: np.ndarray,
                     historical_values: np.ndarray,
                     forecast_times: np.ndarray,
                     forecast_values: np.ndarray,
                     confidence_interval: Optional[Tuple[np.ndarray, np.ndarray]] = None,
                     title: str = "Time Series Forecast") -> Optional[plt.Figure]:
        """
        Plot historical data with forecast.
        
        Args:
            historical_times: Historical time points
            historical_values: Historical values
            forecast_times: Forecast time points
            forecast_values: Forecast values
            confidence_interval: Optional (lower, upper) bounds
            title: Plot title
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        fig, ax = plt.subplots(figsize=(12, 6))
        
        # Historical data
        ax.plot(historical_times, historical_values, 'b-', linewidth=1.5, 
               label='Historical', alpha=0.8)
        
        # Forecast
        ax.plot(forecast_times, forecast_values, 'r--', linewidth=2, 
               label='Forecast')
        
        # Confidence interval
        if confidence_interval is not None:
            lower, upper = confidence_interval
            ax.fill_between(forecast_times, lower, upper, 
                          color='red', alpha=0.2, label='95% Confidence')
        
        # Vertical line at forecast start
        ax.axvline(x=historical_times[-1], color='gray', linestyle='--', 
                  linewidth=1, alpha=0.7)
        
        ax.set_xlabel('Time')
        ax.set_ylabel('Value')
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
        ax.legend(loc='upper left')
        
        plt.tight_layout()
        
        self.plots.append(fig)
        return fig

    def plot_multipanel_dashboard(self,
                                  data: Dict[str, Dict],
                                  title: str = "Monitoring Dashboard") -> Optional[go.Figure]:
        """
        Create multipanel dashboard with Plotly.
        
        Args:
            data: Dictionary of panel name to data
            title: Dashboard title
            
        Returns:
            Plotly figure or None
        """
        if not PLOTLY_AVAILABLE:
            return None
        
        n_panels = len(data)
        fig = make_subplots(
            rows=n_panels, cols=1,
            subplot_titles=list(data.keys()),
            shared_xaxes=True
        )
        
        for i, (name, panel_data) in enumerate(data.items(), 1):
            times = panel_data.get('times', [])
            values = panel_data.get('values', [])
            
            fig.add_trace(
                go.Scatter(x=times, y=values, mode='lines', name=name),
                row=i, col=1
            )
            
            # Add horizontal lines if specified
            if 'thresholds' in panel_data:
                for thresh_name, thresh_value in panel_data['thresholds'].items():
                    fig.add_hline(
                        y=thresh_value,
                        line_dash="dash",
                        line_color="red",
                        opacity=0.5,
                        row=i, col=1
                    )
        
        fig.update_layout(
            title=title,
            height=300 * n_panels,
            showlegend=False
        )
        
        return fig

    def save_plot(self, filename: str, fig=None) -> None:
        """
        Save plot to file.
        
        Args:
            filename: Output filename
            fig: Figure to save (saves last if None)
        """
        if not MATPLOTLIB_AVAILABLE:
            return
        
        if fig is None and self.plots:
            fig = self.plots[-1]
        
        if fig is not None:
            fig.savefig(filename, dpi=300, bbox_inches='tight')
            logger.info(f"Plot saved to {filename}")


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Time Series Plots")
    print("="*60)
    
    tsp = TimeSeriesPlots()
    
    # Generate sample data
    t = np.linspace(0, 10, 100)
    V = 2.5 + 0.2 * np.sin(2 * np.pi * 0.2 * t) + 0.05 * np.random.randn(len(t))
    L = 3.0 + 0.3 * np.cos(2 * np.pi * 0.2 * t) + 0.05 * np.random.randn(len(t))
    F = 0.08 + 0.01 * np.sin(2 * np.pi * 0.1 * t) + 0.002 * np.random.randn(len(t))
    
    if MATPLOTLIB_AVAILABLE:
        # Plot VLF series
        fig = tsp.plot_vlf_series(t, V, L, F, title="Sample VLF Time Series")
        print("\nCreated VLF time series plot")
        
        # Calculate derivatives
        dVdt = np.gradient(V, t)
        dLdt = np.gradient(L, t)
        dFdt = np.gradient(F, t)
        
        # Plot derivatives
        fig2 = tsp.plot_derivatives(t, dVdt, dLdt, dFdt)
        print("Created derivatives plot")
        
        # Correlation matrix
        data_dict = {'V': V, 'L': L, 'F': F}
        fig3 = tsp.plot_correlation_matrix(data_dict)
        print("Created correlation matrix")
    
    print("\n" + "="*60)
