"""
Monthly MODFLOW Results Visualizer
==================================
Creates compelling visualizations for LinkedIn/blog posts showing:
- Seasonal groundwater responses
- Surface-groundwater interactions  
- Water management trade-offs
"""

import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime, timedelta
import seaborn as sns

# Set style for professional plots
plt.style.use('seaborn-v0_8-whitegrid')
sns.set_palette("husl")

class MonthlyResultsVisualizer:
    """Create publication-ready plots from monthly MODFLOW results"""
    
    def __init__(self, results_file='monthly_results.json'):
        self.results_file = results_file
        self.results = self._load_results()
        
    def _load_results(self):
        """Load monthly results from JSON file"""
        try:
            with open(self.results_file, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            print(f"Results file {self.results_file} not found. Run simulation first.")
            return []
    
    def create_dashboard(self, save_path='monthly_dashboard.png'):
        """Create comprehensive dashboard for LinkedIn/blog"""
        
        if not self.results:
            print("No results to visualize")
            return
            
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle('Monthly Surface-Groundwater Management Dashboard', 
                    fontsize=16, fontweight='bold')
        
        # Extract data
        months = [r['month_name'] for r in self.results]
        return_flows = [r['outputs']['avg_return_flow_m3s'] for r in self.results]
        center_heads = [r['outputs']['end_month_head_center_m'] for r in self.results]
        diversions = [r['inputs']['surface_diversion_m3d'] for r in self.results]
        pumping = [r['inputs']['pumping_rate_m3d'] for r in self.results]
        
        # Plot 1: Return Flows vs Surface Diversions
        ax1 = axes[0, 0]
        ax1_twin = ax1.twinx()
        
        line1 = ax1.plot(months, return_flows, 'o-', color='blue', linewidth=2, 
                        markersize=6, label='Return Flow to River')
        line2 = ax1_twin.plot(months, [d/1000 for d in diversions], 's-', 
                             color='red', linewidth=2, markersize=6, 
                             label='Surface Diversion')
        
        ax1.set_ylabel('Return Flow (m³/s)', color='blue', fontweight='bold')
        ax1_twin.set_ylabel('Surface Diversion (×1000 m³/d)', color='red', fontweight='bold')
        ax1.set_title('Surface-Groundwater Interaction', fontweight='bold')
        ax1.tick_params(axis='x', rotation=45)
        
        # Combined legend
        lines = line1 + line2
        labels = [l.get_label() for l in lines]
        ax1.legend(lines, labels, loc='upper left')
        
        # Plot 2: Groundwater Levels
        ax2 = axes[0, 1]
        ax2.plot(months, center_heads, 'o-', color='green', linewidth=3, 
                markersize=8, label='Center Well Head')
        ax2.axhline(y=35, color='red', linestyle='--', alpha=0.7, 
                   label='Critical Level')
        ax2.set_ylabel('Groundwater Head (m)', fontweight='bold')
        ax2.set_title('Groundwater Level Monitoring', fontweight='bold')
        ax2.legend()
        ax2.tick_params(axis='x', rotation=45)
        ax2.grid(True, alpha=0.3)
        
        # Plot 3: Water Budget Components
        ax3 = axes[1, 0]
        width = 0.35
        x = np.arange(len(months))
        
        pumping_scaled = [p/1000 for p in pumping]  # Convert to thousands
        diversions_scaled = [d/1000 for d in diversions]
        
        bars1 = ax3.bar(x - width/2, pumping_scaled, width, label='Groundwater Pumping', 
                       color='orange', alpha=0.8)
        bars2 = ax3.bar(x + width/2, diversions_scaled, width, label='Surface Diversion', 
                       color='lightblue', alpha=0.8)
        
        ax3.set_ylabel('Water Use (×1000 m³/d)', fontweight='bold')
        ax3.set_title('Monthly Water Management', fontweight='bold')
        ax3.set_xticks(x)
        ax3.set_xticklabels(months, rotation=45)
        ax3.legend()
        
        # Add value labels on bars
        for bar in bars1:
            height = bar.get_height()
            ax3.text(bar.get_x() + bar.get_width()/2., height + 0.1,
                    f'{height:.1f}', ha='center', va='bottom', fontsize=8)
        
        # Plot 4: System Efficiency Metrics
        ax4 = axes[1, 1]
        
        # Calculate efficiency metrics
        total_supply = [p + d for p, d in zip(pumping, diversions)]
        gw_fraction = [p/t * 100 if t > 0 else 0 for p, t in zip(pumping, total_supply)]
        
        ax4.plot(months, gw_fraction, 'o-', color='purple', linewidth=2, 
                markersize=6, label='Groundwater Dependency')
        ax4.axhline(y=50, color='gray', linestyle='--', alpha=0.7, 
                   label='50% Target')
        ax4.set_ylabel('Groundwater Dependency (%)', fontweight='bold')
        ax4.set_title('Supply Portfolio Balance', fontweight='bold')
        ax4.legend()
        ax4.tick_params(axis='x', rotation=45)
        ax4.set_ylim(0, 100)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Dashboard saved to: {save_path}")
        plt.show()
    
    def create_story_plot(self, save_path='management_story.png'):
        """Create narrative-focused plot for blog post"""
        
        if not self.results:
            return
            
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
        fig.suptitle('Integrated Water Management: Balancing Surface and Groundwater Resources', 
                    fontsize=14, fontweight='bold')
        
        months = [r['month_name'] for r in self.results]
        return_flows = [r['outputs']['avg_return_flow_m3s'] for r in self.results]
        center_heads = [r['outputs']['end_month_head_center_m'] for r in self.results]
        diversions = [r['inputs']['surface_diversion_m3d']/1000 for r in self.results]
        pumping = [r['inputs']['pumping_rate_m3d']/1000 for r in self.results]
        
        # Top plot: The Management Challenge
        ax1_twin = ax1.twinx()
        
        # Surface water diversions (demand)
        bars = ax1.bar(months, diversions, alpha=0.6, color='lightcoral', 
                      label='Surface Water Diversion')
        
        # Groundwater return flows (system response)
        line = ax1_twin.plot(months, return_flows, 'o-', color='darkblue', 
                           linewidth=3, markersize=8, label='Groundwater Return Flow')
        
        ax1.set_ylabel('Surface Diversion (×1000 m³/d)', fontweight='bold', color='darkred')
        ax1_twin.set_ylabel('Return Flow to River (m³/s)', fontweight='bold', color='darkblue')
        ax1.set_title('The Challenge: Higher Diversions Reduce Groundwater Return Flows', 
                     fontweight='bold')
        
        # Add annotations for key insights
        if len(months) >= 3:
            ax1.annotate('Peak demand period', 
                        xy=(2, max(diversions)), xytext=(1, max(diversions) + 2),
                        arrowprops=dict(arrowstyle='->', color='red'),
                        fontsize=10, ha='center')
        
        ax1.tick_params(axis='x', rotation=45)
        
        # Bottom plot: System Response
        ax2.fill_between(months, center_heads, alpha=0.3, color='green', 
                        label='Groundwater Level')
        ax2.plot(months, center_heads, 'o-', color='darkgreen', linewidth=2, 
                markersize=6)
        
        # Add critical thresholds
        ax2.axhline(y=35, color='red', linestyle='--', linewidth=2, 
                   label='Critical Level (35m)')
        ax2.axhline(y=40, color='orange', linestyle='--', linewidth=2, 
                   label='Management Trigger (40m)')
        
        ax2.set_ylabel('Groundwater Head (m)', fontweight='bold')
        ax2.set_title('System Response: Groundwater Levels Track Seasonal Demand', 
                     fontweight='bold')
        ax2.legend(loc='upper right')
        ax2.tick_params(axis='x', rotation=45)
        ax2.grid(True, alpha=0.3)
        
        # Add insight text box
        insight_text = ("Key Insight: Monthly MODFLOW coupling enables proactive management\n"
                       "by predicting groundwater responses to surface water decisions")
        ax2.text(0.02, 0.98, insight_text, transform=ax2.transAxes, 
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow", alpha=0.8),
                verticalalignment='top', fontsize=10, style='italic')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Story plot saved to: {save_path}")
        plt.show()
    
    def create_technical_summary(self):
        """Print technical summary for blog post"""
        
        if not self.results:
            return
            
        print("\n" + "="*60)
        print("MONTHLY MODFLOW-GOLDSIM COUPLING - TECHNICAL SUMMARY")
        print("="*60)
        
        # Calculate summary statistics
        return_flows = [r['outputs']['avg_return_flow_m3s'] for r in self.results]
        heads = [r['outputs']['end_month_head_center_m'] for r in self.results]
        diversions = [r['inputs']['surface_diversion_m3d'] for r in self.results]
        pumping = [r['inputs']['pumping_rate_m3d'] for r in self.results]
        
        print(f"Simulation Period: {len(self.results)} months")
        print(f"Model Domain: 5km × 5km (25×25 cells, 200m spacing)")
        print(f"Temporal Resolution: 30-day stress periods, daily time steps")
        
        print(f"\nReturn Flow Statistics:")
        print(f"  Range: {min(return_flows):.3f} - {max(return_flows):.3f} m³/s")
        print(f"  Average: {np.mean(return_flows):.3f} m³/s")
        print(f"  Seasonal variation: {(max(return_flows)-min(return_flows))/np.mean(return_flows)*100:.1f}%")
        
        print(f"\nGroundwater Level Statistics:")
        print(f"  Range: {min(heads):.1f} - {max(heads):.1f} m")
        print(f"  Average: {np.mean(heads):.1f} m")
        print(f"  Seasonal drawdown: {max(heads)-min(heads):.1f} m")
        
        print(f"\nWater Management Statistics:")
        print(f"  Surface diversions: {min(diversions)/1000:.1f} - {max(diversions)/1000:.1f} thousand m³/d")
        print(f"  Groundwater pumping: {min(pumping)/1000:.1f} - {max(pumping)/1000:.1f} thousand m³/d")
        
        # Calculate correlation
        correlation = np.corrcoef(diversions, return_flows)[0, 1]
        print(f"\nSurface-Groundwater Coupling:")
        print(f"  Correlation (diversions vs return flows): {correlation:.3f}")
        print(f"  Interpretation: {'Strong negative' if correlation < -0.7 else 'Moderate negative' if correlation < -0.3 else 'Weak'} coupling")
        
        print("\n" + "="*60)


def create_sample_results():
    """Create sample results for demonstration"""
    
    # Realistic seasonal water management scenario
    seasonal_patterns = [
        # Month, Surface Div, Pumping, Recharge, Expected Return Flow, Expected Head
        ("Jan", 5000, 2000, 0.001, 0.15, 42.5),
        ("Feb", 4000, 1800, 0.002, 0.18, 42.8),
        ("Mar", 6000, 2200, 0.003, 0.12, 42.3),
        ("Apr", 8000, 2800, 0.002, 0.08, 41.8),
        ("May", 12000, 3500, 0.001, 0.04, 41.2),
        ("Jun", 15000, 4000, 0.0005, 0.02, 40.5),
    ]
    
    results = []
    for i, (month, div, pump, rech, ret_flow, head) in enumerate(seasonal_patterns):
        result = {
            'month': i + 1,
            'month_name': month,
            'year': 2024,
            'inputs': {
                'surface_diversion_m3d': div,
                'pumping_rate_m3d': pump,
                'recharge_rate_md': rech
            },
            'outputs': {
                'avg_return_flow_m3s': ret_flow,
                'avg_return_flow_m3d': ret_flow * 86400,
                'end_month_head_center_m': head,
                'end_month_head_avg_m': head - 1.0,
                'storage_change_m3': -pump * 30 * 0.8,
                'water_balance_error_pct': 0.5
            },
            'timestamp': datetime.now().isoformat()
        }
        results.append(result)
    
    # Save sample results
    with open('monthly_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    
    print("Sample results created: monthly_results.json")
    return results


if __name__ == "__main__":
    print("Monthly MODFLOW Results Visualizer")
    print("==================================")
    
    # Create sample data if no results exist
    try:
        with open('monthly_results.json', 'r') as f:
            results = json.load(f)
        print(f"Loaded {len(results)} months of results")
    except FileNotFoundError:
        print("No results found. Creating sample data...")
        results = create_sample_results()
    
    # Create visualizations
    viz = MonthlyResultsVisualizer()
    
    print("\nCreating dashboard...")
    viz.create_dashboard()
    
    print("\nCreating story plot...")
    viz.create_story_plot()
    
    print("\nGenerating technical summary...")
    viz.create_technical_summary()
    
    print("\nVisualization complete! Ready for LinkedIn/blog post.")