# part2_visualization.py

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

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"""
        fig = make_subplots(
            rows=4, cols=1,  # Added one more row for NDVI
            subplot_titles=(
                '<b>Temperature (°C)</b>',
                '<b>Humidity (%)</b>',
                '<b>Rainfall (mm/day)</b>',
                '<b>Vegetation Index (NDVI)</b>'
            ),
            vertical_spacing=0.08,
            row_heights=[0.25, 0.25, 0.25, 0.25]
        )
        
        # Add standard weather plots
        self.add_weather_plots(fig, df)
        
        # Add NDVI plot
        self.add_ndvi_plot(fig, df)
        
        # Update layout
        fig.update_layout(
            height=1000,  # Increased height for additional plot
            showlegend=True,
            title={
                'text': "Enhanced Tobacco Growing 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)
        )

        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_weather_plots(self, fig, df):
        """Add weather-related plots"""
        # Temperature plot
        for data_type, color in [('historical', 'royalblue'), 
                               ('forecast_5day', 'firebrick'),
                               ('forecast_extended', 'rgba(255, 165, 0, 0.5)')]:
            mask = df['type'] == data_type
            if any(mask):
                fig.add_trace(
                    go.Scatter(
                        x=df[mask]['date'],
                        y=df[mask]['temperature'],
                        name=f'{data_type.replace("_", " ").title()} Temperature',
                        line=dict(color=color, width=2),
                        mode='lines'
                    ),
                    row=1, col=1
                )

        # Add temperature rolling average
        fig.add_trace(
            go.Scatter(
                x=df['date'],
                y=df['temp_7day_avg'],
                name='7-day Temperature Average',
                line=dict(color='purple', width=1, dash='dot'),
                mode='lines'
            ),
            row=1, col=1
        )

        # Humidity plot
        for data_type, color in [('historical', 'royalblue'), 
                               ('forecast_5day', 'firebrick'),
                               ('forecast_extended', 'rgba(255, 165, 0, 0.5)')]:
            mask = df['type'] == data_type
            if any(mask):
                fig.add_trace(
                    go.Scatter(
                        x=df[mask]['date'],
                        y=df[mask]['humidity'],
                        name=f'{data_type.replace("_", " ").title()} Humidity',
                        line=dict(color=color, width=2),
                        mode='lines'
                    ),
                    row=2, col=1
                )

        # Add humidity rolling average
        fig.add_trace(
            go.Scatter(
                x=df['date'],
                y=df['humidity_7day_avg'],
                name='7-day Humidity Average',
                line=dict(color='purple', width=1, dash='dot'),
                mode='lines'
            ),
            row=2, col=1
        )

        # Rainfall plot
        for data_type, color in [('historical', 'royalblue'), 
                               ('forecast_5day', 'firebrick'),
                               ('forecast_extended', 'rgba(255, 165, 0, 0.5)')]:
            mask = df['type'] == data_type
            if any(mask):
                fig.add_trace(
                    go.Bar(
                        x=df[mask]['date'],
                        y=df[mask]['rainfall'],
                        name=f'{data_type.replace("_", " ").title()} Rainfall',
                        marker_color=color
                    ),
                    row=3, col=1
                )

    def add_ndvi_plot(self, fig, df):
        """Add NDVI plot"""
        # Historical NDVI
        mask = df['type'] == 'historical'
        if any(mask):
            fig.add_trace(
                go.Scatter(
                    x=df[mask]['date'],
                    y=df[mask]['estimated_ndvi'],
                    name='Historical NDVI',
                    line=dict(color='green', width=2),
                    mode='lines'
                ),
                row=4, col=1
            )

        # Forecast NDVI
        mask = df['type'].isin(['forecast_5day', 'forecast_extended'])
        if any(mask):
            fig.add_trace(
                go.Scatter(
                    x=df[mask]['date'],
                    y=df[mask]['estimated_ndvi'],
                    name='Forecast NDVI',
                    line=dict(color='orange', width=2, dash='dot'),
                    mode='lines'
                ),
                row=4, col=1
            )

        # Add optimal NDVI range
        fig.add_hline(y=self.optimal_conditions['ndvi']['min'], 
                     line_dash="dash", line_color="green",
                     annotation_text="Min Optimal NDVI", 
                     row=4, col=1)
        fig.add_hline(y=self.optimal_conditions['ndvi']['max'], 
                     line_dash="dash", line_color="green",
                     annotation_text="Max Optimal NDVI", 
                     row=4, col=1)

    def create_enhanced_map(self, lat, lon, score, ndvi_value):
        """Create an interactive map with both weather and vegetation analysis"""
        m = folium.Map(location=[lat, lon], zoom_start=13)
        
        # Add base 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='NDVI Values'
        )
        
        # Add NDVI-based circle
        ndvi_color = ndvi_colormap(ndvi_value)
        folium.Circle(
            radius=2000,
            location=[lat, lon],
            popup=f'NDVI: {ndvi_value:.2f}',
            color=ndvi_color,
            fill=True,
            fillOpacity=0.4
        ).add_to(m)
        
        # Add score-based circles
        score_color = self.get_score_color(score)
        for radius in [500, 1000, 1500]:
            folium.Circle(
                radius=radius,
                location=[lat, lon],
                color=score_color,
                popup=f'Growing Score: {score:.2f}',
                fill=False,
                weight=2
            ).add_to(m)
        
        # Add measurement tools
        plugins.MeasureControl(position='topright').add_to(m)
        
        # Add fullscreen option
        plugins.Fullscreen().add_to(m)
        
        # Add mini map
        minimap = plugins.MiniMap()
        m.add_child(minimap)
        
        # Add layer control
        folium.LayerControl().add_to(m)
        
        # Add colormap to map
        m.add_child(ndvi_colormap)
        
        return m._repr_html_()

    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'

    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