import plotly.graph_objects as go
from plotly.subplots import make_subplots
import folium
from folium import plugins
import numpy as np
import branca.colormap as cm
from datetime import datetime

class VisualizationHandler:
    def __init__(self, optimal_conditions):
        self.optimal_conditions = optimal_conditions
        self.ndvi_colors = [
            '#d73027',  # Very low vegetation
            '#f46d43',  # Low vegetation
            '#fdae61',  # Sparse vegetation
            '#fee08b',  # Moderate vegetation
            '#d9ef8b',  # Good vegetation
            '#a6d96a',  # High vegetation
            '#66bd63',  # Very high vegetation
            '#1a9850'   # Dense vegetation
        ]

    def create_interactive_plots(self, df):
        """Create enhanced interactive Plotly visualizations"""
        if df is None or df.empty:
            return go.Figure()  # Return empty figure if no data
            
        fig = make_subplots(
            rows=4, cols=1,
            subplot_titles=(
                '<b>Temperature Pattern (°C)</b>',
                '<b>Humidity Pattern (%)</b>',
                '<b>Rainfall Pattern (mm/day)</b>',
                '<b>Vegetation & Suitability Indices</b>'
            ),
            vertical_spacing=0.08,
            row_heights=[0.28, 0.24, 0.24, 0.24]
        )

        # Add temperature visualization
        self.add_temperature_plot(fig, df)
        
        # Add humidity visualization
        self.add_humidity_plot(fig, df)
        
        # Add rainfall visualization
        self.add_rainfall_plot(fig, df)
        
        # Add vegetation and suitability visualization
        self.add_combined_indices_plot(fig, df)

        # Update layout
        fig.update_layout(
            height=1000,
            showlegend=True,
            title={
                'text': "Agricultural Conditions Analysis",
                'y':0.95,
                'x':0.5,
                'xanchor': 'center',
                'yanchor': 'top',
                'font': dict(size=20)
            },
            paper_bgcolor='white',
            plot_bgcolor='rgba(0,0,0,0.05)',
            font=dict(size=12),
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            margin=dict(l=60, r=30, t=100, b=60)
        )

        # Add season shading
        self.add_season_shading(fig, df)
        
        # Update axes for all subplots
        fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='rgba(0,0,0,0.1)')
        fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='rgba(0,0,0,0.1)')
        
        return fig

    def add_temperature_plot(self, fig, df):
        """Add temperature visualization with range"""
        # Temperature range area
        if 'temp_max' in df.columns and 'temp_min' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['temp_max'],
                    name='Max Temperature',
                    line=dict(color='rgba(255,0,0,0.0)'),
                    showlegend=False
                ),
                row=1, col=1
            )
            
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['temp_min'],
                    name='Temperature Range',
                    fill='tonexty',
                    fillcolor='rgba(255,0,0,0.1)',
                    line=dict(color='rgba(255,0,0,0.0)')
                ),
                row=1, col=1
            )

        # Main temperature line
        fig.add_trace(
            go.Scatter(
                x=df['date'],
                y=df['temperature'],
                name='Temperature',
                line=dict(color='red', width=2),
                mode='lines'
            ),
            row=1, col=1
        )

        # Add rolling average
        if 'temp_7day_avg' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['temp_7day_avg'],
                    name='7-Day Average',
                    line=dict(color='darkred', width=1, dash='dot'),
                    mode='lines'
                ),
                row=1, col=1
            )

        # Add optimal range
        for limit_type, value in self.optimal_conditions['temperature'].items():
            fig.add_hline(
                y=value,
                line_dash="dash",
                line_color="green",
                annotation_text=f"Optimal {limit_type}",
                row=1, col=1
            )

    def add_humidity_plot(self, fig, df):
        """Add humidity visualization"""
        fig.add_trace(
            go.Scatter(
                x=df['date'],
                y=df['humidity'],
                name='Humidity',
                line=dict(color='blue', width=2),
                mode='lines'
            ),
            row=2, col=1
        )

        if 'humidity_7day_avg' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['humidity_7day_avg'],
                    name='7-Day Average',
                    line=dict(color='darkblue', width=1, dash='dot'),
                    mode='lines'
                ),
                row=2, col=1
            )

        # Add optimal range
        for limit_type, value in self.optimal_conditions['humidity'].items():
            fig.add_hline(
                y=value,
                line_dash="dash",
                line_color="green",
                annotation_text=f"Optimal {limit_type}",
                row=2, col=1
            )

    def add_rainfall_plot(self, fig, df):
        """Add rainfall visualization"""
        fig.add_trace(
            go.Bar(
                x=df['date'],
                y=df['rainfall'],
                name='Daily Rainfall',
                marker_color='lightblue',
                opacity=0.6
            ),
            row=3, col=1
        )

        if 'rainfall_7day_avg' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['rainfall_7day_avg'],
                    name='7-Day Average',
                    line=dict(color='blue', width=2),
                    mode='lines'
                ),
                row=3, col=1
            )

    def add_combined_indices_plot(self, fig, df):
        """Add vegetation and suitability indices visualization"""
        if 'estimated_ndvi' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['estimated_ndvi'],
                    name='Vegetation Index',
                    line=dict(color='green', width=2),
                    mode='lines'
                ),
                row=4, col=1
            )

        if 'daily_suitability' in df.columns:
            fig.add_trace(
                go.Scatter(
                    x=df['date'],
                    y=df['daily_suitability'],
                    name='Growing Suitability',
                    line=dict(color='purple', width=2),
                    mode='lines'
                ),
                row=4, col=1
            )

    def add_season_shading(self, fig, df):
        """Add season shading to all plots"""
        if 'season' in df.columns:
            seasons = df['season'].unique()
            season_colors = {
                'Main': 'rgba(0,255,0,0.1)',    # Green
                'Early': 'rgba(255,255,0,0.1)',  # Yellow
                'Late': 'rgba(255,165,0,0.1)',   # Orange
                'Dry': 'rgba(255,0,0,0.1)'       # Red
            }
            
            for season in seasons:
                season_data = df[df['season'] == season]
                if not season_data.empty:
                    for row in range(1, 5):
                        fig.add_vrect(
                            x0=season_data['date'].iloc[0],
                            x1=season_data['date'].iloc[-1],
                            fillcolor=season_colors.get(season, 'rgba(128,128,128,0.1)'),
                            layer="below",
                            line_width=0,
                            annotation_text=season if row == 1 else None,
                            annotation_position="top left",
                            row=row, col=1
                        )

    def create_enhanced_map(self, lat, lon, score, ndvi_value):
        """Create an interactive map with analysis overlays"""
        m = folium.Map(location=[lat, lon], zoom_start=13)
        
        # Add measurement tools
        plugins.MeasureControl(position='topright').add_to(m)
        plugins.Fullscreen().add_to(m)
        
        # Add location marker
        folium.Marker(
            [lat, lon],
            popup='Analysis Location',
            icon=folium.Icon(color='red', icon='info-sign')
        ).add_to(m)
        
        # Create NDVI colormap
        ndvi_colormap = cm.LinearColormap(
            colors=self.ndvi_colors,
            vmin=-1,
            vmax=1,
            caption='Vegetation Index (NDVI)'
        )
        
        # Add NDVI circle
        folium.Circle(
            radius=2000,
            location=[lat, lon],
            popup=f'NDVI: {ndvi_value:.2f}',
            color=ndvi_colormap(ndvi_value),
            fill=True,
            fillOpacity=0.4
        ).add_to(m)
        
        # Add suitability circles
        score_color = self.get_score_color(score)
        for radius in [500, 1000, 1500]:
            folium.Circle(
                radius=radius,
                location=[lat, lon],
                popup=f'Suitability Score: {score:.2f}',
                color=score_color,
                fill=False,
                weight=2
            ).add_to(m)
        
        # Add mini map
        minimap = plugins.MiniMap()
        m.add_child(minimap)
        
        # Add layer control
        folium.LayerControl().add_to(m)
        m.add_child(ndvi_colormap)
        
        return m._repr_html_()

    def create_gauge_chart(self, score):
        """Create an enhanced gauge chart for the overall score"""
        fig = go.Figure(go.Indicator(
            mode="gauge+number+delta",
            value=score,
            domain={'x': [0, 1], 'y': [0, 1]},
            title={
                'text': "Growing Conditions Score",
                'font': {'size': 24}
            },
            delta={
                'reference': 0.8,
                'increasing': {'color': "green"},
                'decreasing': {'color': "red"}
            },
            gauge={
                'axis': {'range': [None, 1], 'tickwidth': 1, 'tickcolor': "darkblue"},
                'bar': {'color': "darkblue"},
                'bgcolor': "white",
                'borderwidth': 2,
                'bordercolor': "gray",
                'steps': [
                    {'range': [0, 0.4], 'color': 'rgba(255, 0, 0, 0.6)'},
                    {'range': [0.4, 0.6], 'color': 'rgba(255, 255, 0, 0.6)'},
                    {'range': [0.6, 0.8], 'color': 'rgba(144, 238, 144, 0.6)'},
                    {'range': [0.8, 1], 'color': 'rgba(0, 128, 0, 0.6)'}
                ],
                'threshold': {
                    'line': {'color': "red", 'width': 4},
                    'thickness': 0.75,
                    'value': 0.8
                }
            }
        ))
        
        fig.update_layout(
            height=300,
            margin=dict(l=20, r=20, t=60, b=20),
            paper_bgcolor="white",
            font={'color': "darkblue", 'family': "Arial"}
        )
        
        return fig

    def get_score_color(self, score):
        """Get color based on score"""
        if score >= 0.8:
            return 'green'
        elif score >= 0.6:
            return 'yellow'
        elif score >= 0.4:
            return 'orange'
        return 'red'