"""
MODFLOW Tutorial 2 Style Setup
==============================
Clear demonstration of groundwater pumping effects

Key Setup (like Tutorial 2):
- Confined aquifer (constant properties)
- Constant head boundaries (40m on all edges)
- Single pumping well in center
- Steady-state solution
- Clear radial drawdown pattern
"""

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

def explain_setup():
    """Explain the model setup clearly"""
    print("MODFLOW TUTORIAL 2 STYLE SETUP")
    print("=" * 40)
    print("Domain: 2.1 km x 2.1 km (21x21 cells, 100m each)")
    print("Aquifer: Confined, 50m thick")
    print("Boundaries: Constant head (40m) on all edges")
    print("Well: Center of domain (1050m, 1050m)")
    print("Properties:")
    print("  - Hydraulic conductivity: 10 m/day")
    print("  - Transmissivity: 500 m²/day (K × thickness)")
    print("  - Specific storage: 1e-5 /m")
    print("Solution: Steady-state (equilibrium)")
    print()

def run_single_scenario(pumping_rate, description):
    """Run single pumping scenario with detailed output"""
    
    print(f"SCENARIO: {description}")
    print(f"Pumping rate: {pumping_rate} m³/day")
    print("-" * 30)
    
    # Model setup
    nlay, nrow, ncol = 1, 21, 21
    delr = delc = 100.0
    top, botm = 50.0, 0.0
    
    model_ws = 'tutorial_demo'
    if not os.path.exists(model_ws):
        os.makedirs(model_ws)
    
    mf = flopy.modflow.Modflow(modelname='tutorial', model_ws=model_ws)
    
    # Packages
    dis = flopy.modflow.ModflowDis(mf, nlay, nrow, ncol, delr=delr, delc=delc,
                                   top=top, botm=botm, nper=1, perlen=[1.0], 
                                   nstp=[1], steady=[True])
    
    # Boundaries and initial conditions
    ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
    ibound[0, 0, :] = -1   # North: constant head
    ibound[0, -1, :] = -1  # South: constant head
    ibound[0, :, 0] = -1   # West: constant head
    ibound[0, :, -1] = -1  # East: constant head
    
    strt = np.full((nlay, nrow, ncol), 40.0, dtype=np.float32)
    bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)
    
    # Aquifer properties
    hk = 10.0      # Hydraulic conductivity
    transmissivity = hk * (top - botm)  # T = K × b
    print(f"Transmissivity: {transmissivity} m²/day")
    
    lpf = flopy.modflow.ModflowLpf(mf, hk=hk, vka=hk, sy=0.0, ss=1e-5, laytyp=0)
    
    # Well (if pumping)
    if pumping_rate > 0:
        well_row, well_col = 10, 10  # Center
        wel_data = [[0, well_row, well_col, -pumping_rate]]
        wel = flopy.modflow.ModflowWel(mf, stress_period_data={0: wel_data})
        print(f"Well location: Row {well_row+1}, Col {well_col+1} (center)")
    
    # Output and solver
    oc = flopy.modflow.ModflowOc(mf, stress_period_data={(0, 0): ['save head']})
    pcg = flopy.modflow.ModflowPcg(mf, mxiter=500, hclose=1e-6)
    
    # Run model
    mf.write_input()
    success, buff = mf.run_model(silent=True)
    
    if not success:
        print("ERROR: Model failed to converge")
        return None
    
    # Read and analyze results
    hds_file = os.path.join(model_ws, f'{mf.name}.hds')
    hds = flopy.utils.HeadFile(hds_file)
    heads = hds.get_data()[0]  # Remove layer dimension
    
    # Results
    center_head = heads[10, 10]
    boundary_head = heads[0, 10]  # North boundary
    drawdown = 40.0 - center_head if pumping_rate > 0 else 0.0
    
    print(f"Results:")
    print(f"  Center head: {center_head:.3f} m")
    print(f"  Boundary head: {boundary_head:.3f} m")
    if pumping_rate > 0:
        print(f"  Drawdown: {drawdown:.3f} m")
        print(f"  Radius of influence: ~{estimate_radius_of_influence(heads):.0f} m")
    print()
    
    return {
        'heads': heads,
        'center_head': center_head,
        'drawdown': drawdown,
        'pumping_rate': pumping_rate,
        'description': description
    }

def estimate_radius_of_influence(heads):
    """Estimate radius where drawdown becomes negligible"""
    center_head = heads[10, 10]
    boundary_head = 40.0
    
    # Find where head is within 1% of boundary head
    threshold = boundary_head - 0.01 * (boundary_head - center_head)
    
    # Check along east-west line through center
    for j in range(10, 21):
        if heads[10, j] >= threshold:
            return (j - 10) * 100  # Distance in meters
    
    return 1000  # Default if not found

def create_detailed_plot(result):
    """Create detailed plot for single scenario"""
    
    if result is None:
        return
    
    heads = result['heads']
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # Contour plot
    x = np.arange(21) * 100
    y = np.arange(21) * 100
    X, Y = np.meshgrid(x, y)
    
    if result['pumping_rate'] == 0:
        levels = [39.99, 40.0, 40.01]
    else:
        min_head = np.min(heads)
        levels = np.linspace(min_head, 40.0, 15)
    
    cs = ax1.contour(X, Y, heads, levels=levels, colors='blue')
    ax1.clabel(cs, inline=True, fontsize=9, fmt='%.2f')
    cf = ax1.contourf(X, Y, heads, levels=levels, cmap='Blues_r', alpha=0.7)
    
    # Mark well and boundaries
    if result['pumping_rate'] > 0:
        ax1.plot(1000, 1000, 'ro', markersize=10, label='Pumping Well')
    
    # Boundary markers
    ax1.plot([0, 2000, 2000, 0, 0], [0, 0, 2000, 2000, 0], 'k-', linewidth=3)
    ax1.text(1000, -100, 'Constant Head = 40m', ha='center', fontweight='bold')
    
    ax1.set_xlabel('Distance (m)')
    ax1.set_ylabel('Distance (m)')
    ax1.set_title(f'{result["description"]}\nHead Contours (m)')
    ax1.set_aspect('equal')
    ax1.grid(True, alpha=0.3)
    if result['pumping_rate'] > 0:
        ax1.legend()
    
    # Cross-section
    distances = (np.arange(21) - 10) * 100
    cross_section = heads[10, :]  # Through center
    
    ax2.plot(distances, cross_section, 'b-', linewidth=3, marker='o', markersize=5)
    ax2.axhline(y=40.0, color='red', linestyle='--', alpha=0.7, label='Boundary Head')
    ax2.axvline(x=0, color='red', linestyle=':', alpha=0.7, label='Well Location')
    
    ax2.set_xlabel('Distance from Well (m)')
    ax2.set_ylabel('Hydraulic Head (m)')
    ax2.set_title('Head Profile (East-West)')
    ax2.grid(True, alpha=0.3)
    ax2.legend()
    
    # Add text box with key info
    info_text = f"Pumping: {result['pumping_rate']} m³/day\n"
    info_text += f"Center Head: {result['center_head']:.3f} m\n"
    if result['pumping_rate'] > 0:
        info_text += f"Drawdown: {result['drawdown']:.3f} m"
    
    ax2.text(0.02, 0.98, info_text, transform=ax2.transAxes, 
             verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat'))
    
    plt.tight_layout()
    plt.savefig(f'tutorial_demo/{result["description"].replace(" ", "_").lower()}.png', 
                dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    explain_setup()
    
    # Test scenarios
    scenarios = [
        (0, "No Pumping (Baseline)"),
        (1000, "Moderate Pumping"),
        (2000, "Heavy Pumping")
    ]
    
    results = []
    for pump_rate, desc in scenarios:
        result = run_single_scenario(pump_rate, desc)
        if result:
            results.append(result)
            create_detailed_plot(result)
    
    # Summary comparison
    if len(results) > 1:
        print("COMPARISON SUMMARY")
        print("=" * 30)
        for result in results:
            print(f"{result['description']:20}: {result['center_head']:.3f} m", end="")
            if result['pumping_rate'] > 0:
                print(f" ({result['drawdown']:.3f} m drawdown)")
            else:
                print(" (baseline)")