"""
Simple Groundwater Pumping Demonstration
========================================
Based on MODFLOW Tutorial 2 - shows clear water table response to pumping

Simple setup:
- Confined aquifer (constant thickness)
- Constant head boundaries (like lakes/rivers)
- Single pumping well in center
- Clear, predictable drawdown pattern
"""

import numpy as np
import flopy
import matplotlib.pyplot as plt
import os

def create_simple_model(pumping_rate=0):
    """
    Create simple MODFLOW model like Tutorial 2
    
    Args:
        pumping_rate: Pumping rate in m³/day (positive = pumping out)
    """
    
    # Simple grid - 21x21 for symmetry
    nlay, nrow, ncol = 1, 21, 21
    delr = delc = 100.0  # 100m cells = 2.1km x 2.1km domain
    
    # Confined aquifer
    top = 50.0      # Top of aquifer
    botm = 0.0      # Bottom of aquifer
    thickness = 50.0  # Constant 50m thick
    
    # Steady state for clear demonstration
    nper = 1
    perlen = [1.0]
    nstp = [1]
    steady = [True]
    
    # Create model
    model_ws = 'simple_demo'
    if not os.path.exists(model_ws):
        os.makedirs(model_ws)
        
    mf = flopy.modflow.Modflow(modelname='simple', model_ws=model_ws)
    
    # DIS package
    dis = flopy.modflow.ModflowDis(mf, nlay, nrow, ncol, delr=delr, delc=delc,
                                   top=top, botm=botm, nper=nper,
                                   perlen=perlen, nstp=nstp, steady=steady)
    
    # BAS package - constant head boundaries like Tutorial 2
    ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
    
    # Constant head boundaries on all edges (like surrounding lakes)
    ibound[0, 0, :] = -1   # North edge
    ibound[0, -1, :] = -1  # South edge  
    ibound[0, :, 0] = -1   # West edge
    ibound[0, :, -1] = -1  # East edge
    
    # Initial heads - simple uniform field
    strt = np.full((nlay, nrow, ncol), 40.0, dtype=np.float32)
    
    bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)
    
    # LPF package - confined aquifer (Tutorial 2 style)
    hk = 10.0     # Hydraulic conductivity (m/day)
    vka = 10.0    # Vertical K
    sy = 0.0      # No specific yield (confined)
    ss = 1e-5     # Small specific storage
    laytyp = 0    # Confined aquifer
    
    lpf = flopy.modflow.ModflowLpf(mf, hk=hk, vka=vka, sy=sy, ss=ss,
                                   laytyp=laytyp)
    
    # Add pumping well in center if specified
    if pumping_rate > 0:
        well_row = ncol // 2  # Center row (10)
        well_col = nrow // 2  # Center col (10)
        
        # Negative rate = pumping out
        wel_data = [[0, well_row, well_col, -pumping_rate]]
        wel_spd = {0: wel_data}
        wel = flopy.modflow.ModflowWel(mf, stress_period_data=wel_spd)
    
    # OC package - save heads
    spd = {(0, 0): ['save head']}
    oc = flopy.modflow.ModflowOc(mf, stress_period_data=spd)
    
    # PCG solver
    pcg = flopy.modflow.ModflowPcg(mf, mxiter=500, iter1=50, hclose=1e-6)
    
    return mf

def run_pumping_scenarios():
    """Run scenarios with different pumping rates"""
    
    # Pumping scenarios (m³/day)
    scenarios = [
        (0, "No pumping"),
        (500, "Light pumping (500 m³/day)"),
        (1000, "Moderate pumping (1000 m³/day)"),
        (2000, "Heavy pumping (2000 m³/day)")
    ]
    
    results = {}
    
    for pump_rate, description in scenarios:
        print(f"\nRunning: {description}")
        
        # Create and run model
        mf = create_simple_model(pump_rate)
        mf.write_input()
        success, buff = mf.run_model(silent=True)
        
        if not success:
            print(f"Model failed for {description}")
            continue
            
        # Read results
        hds_file = os.path.join(mf.model_ws, f'{mf.name}.hds')
        hds = flopy.utils.HeadFile(hds_file)
        heads = hds.get_data()
        
        # Store results
        results[pump_rate] = {
            'heads': heads[0],  # Remove layer dimension
            'description': description,
            'center_head': heads[0, 10, 10],  # Center cell
            'max_head': np.max(heads[0]),
            'min_head': np.min(heads[0])
        }
        
        print(f"  Center head: {heads[0, 10, 10]:.2f} m")
        print(f"  Head range: {np.min(heads[0]):.2f} to {np.max(heads[0]):.2f} m")
        if pump_rate > 0:
            drawdown = 40.0 - heads[0, 10, 10]
            print(f"  Drawdown at well: {drawdown:.2f} m")
    
    return results

def plot_results(results):
    """Create clear plots showing pumping effects"""
    
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle('Groundwater Response to Pumping', fontsize=16, fontweight='bold')
    
    pump_rates = [0, 500, 1000, 2000]
    
    for i, pump_rate in enumerate(pump_rates):
        if pump_rate not in results:
            continue
            
        ax = axes[i//2, i%2]
        heads = results[pump_rate]['heads']
        
        # Create contour plot
        x = np.arange(21) * 100  # Convert to meters
        y = np.arange(21) * 100
        X, Y = np.meshgrid(x, y)
        
        # Contour levels
        if pump_rate == 0:
            levels = np.linspace(39.9, 40.1, 11)
        else:
            min_head = np.min(heads)
            levels = np.linspace(min_head, 40.0, 15)
        
        cs = ax.contour(X, Y, heads, levels=levels, colors='blue', linewidths=0.8)
        ax.clabel(cs, inline=True, fontsize=8, fmt='%.1f')
        
        # Fill contours for better visualization
        cf = ax.contourf(X, Y, heads, levels=levels, cmap='Blues_r', alpha=0.6)
        
        # Mark well location
        if pump_rate > 0:
            ax.plot(1000, 1000, 'ro', markersize=8, label='Pumping Well')
            ax.legend()
        
        # Mark boundaries
        ax.plot([0, 2000, 2000, 0, 0], [0, 0, 2000, 2000, 0], 'k-', linewidth=2, 
                label='Constant Head Boundary')
        
        ax.set_xlabel('Distance (m)')
        ax.set_ylabel('Distance (m)')
        ax.set_title(f'{results[pump_rate]["description"]}\nCenter Head: {results[pump_rate]["center_head"]:.2f} m')
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('simple_demo/pumping_effects.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # Create drawdown comparison
    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
    
    pump_rates_plot = []
    center_heads = []
    drawdowns = []
    
    for pump_rate in sorted(results.keys()):
        pump_rates_plot.append(pump_rate)
        center_heads.append(results[pump_rate]['center_head'])
        drawdowns.append(40.0 - results[pump_rate]['center_head'])
    
    ax.plot(pump_rates_plot, drawdowns, 'bo-', linewidth=2, markersize=8)
    ax.set_xlabel('Pumping Rate (m³/day)')
    ax.set_ylabel('Drawdown at Well (m)')
    ax.set_title('Drawdown vs Pumping Rate', fontsize=14, fontweight='bold')
    ax.grid(True, alpha=0.3)
    
    # Add text annotations
    for i, (pump, dd) in enumerate(zip(pump_rates_plot, drawdowns)):
        if pump > 0:
            ax.annotate(f'{dd:.2f} m', (pump, dd), 
                       textcoords="offset points", xytext=(0,10), ha='center')
    
    plt.tight_layout()
    plt.savefig('simple_demo/drawdown_curve.png', dpi=300, bbox_inches='tight')
    plt.show()

def create_cross_section_plot(results):
    """Create cross-section through pumping well"""
    
    fig, ax = plt.subplots(1, 1, figsize=(12, 6))
    
    # Distance from center (well at x=0)
    distances = (np.arange(21) - 10) * 100  # -1000 to +1000 m
    
    colors = ['blue', 'green', 'orange', 'red']
    pump_rates = [0, 500, 1000, 2000]
    
    for i, pump_rate in enumerate(pump_rates):
        if pump_rate not in results:
            continue
            
        heads = results[pump_rate]['heads']
        # Cross-section through center (row 10)
        cross_section = heads[10, :]
        
        ax.plot(distances, cross_section, color=colors[i], linewidth=2, 
                marker='o', markersize=4, label=results[pump_rate]['description'])
    
    ax.axvline(x=0, color='red', linestyle='--', alpha=0.7, label='Pumping Well')
    ax.set_xlabel('Distance from Well (m)')
    ax.set_ylabel('Hydraulic Head (m)')
    ax.set_title('Head Profile Through Pumping Well', fontsize=14, fontweight='bold')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('simple_demo/cross_section.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    print("Simple Groundwater Pumping Demonstration")
    print("=" * 50)
    
    # Run scenarios
    results = run_pumping_scenarios()
    
    # Create plots
    if results:
        print("\nCreating plots...")
        plot_results(results)
        create_cross_section_plot(results)
        print("\nPlots saved in simple_demo/ directory")
        
        # Summary
        print("\nSUMMARY:")
        print("--------")
        for pump_rate in sorted(results.keys()):
            center_head = results[pump_rate]['center_head']
            if pump_rate == 0:
                print(f"No pumping: {center_head:.2f} m (baseline)")
            else:
                drawdown = 40.0 - center_head
                print(f"{pump_rate:4d} m³/day: {center_head:.2f} m ({drawdown:.2f} m drawdown)")
    else:
        print("No results to plot")