"""
Responsive GSPy-MODFLOW Interface
================================
Clean interface for GoldSim-MODFLOW coupling with Tutorial 2 style setup.

Features:
- 21x21 grid (2.1km domain, 100m cells)
- Confined aquifer with constant head boundaries
- 30-day monthly simulations
- Automatic simulation detection
- Monthly persistence within simulations
"""

import os
import time
import numpy as np
import flopy
import pickle
from datetime import datetime

try:
    import gspy
    GSPy_AVAILABLE = True
except ImportError:
    GSPy_AVAILABLE = False

# Global state for simulation detection
_call_counter = 0
_last_call_time = 0

def log_message(message, level=2):
    """Simple logging with fallback"""
    if GSPy_AVAILABLE:
        gspy.log(message, level)
    else:
        print(f"DEBUG: {message}")

def create_modflow_model(well_rate, recharge_rate, well_row, well_col, hydraulic_k, specific_yield):
    """
    Create MODFLOW model with Tutorial 2 style setup
    
    Args:
        well_rate: Pumping rate (m³/day, positive = pumping out)
        recharge_rate: Total recharge (m³/day, distributed over domain)
        well_row: Well row (0-20, center is 10)
        well_col: Well column (0-20, center is 10)
        hydraulic_k: Hydraulic conductivity (m/day)
        specific_yield: Ignored for confined aquifer
    """
    
    # Grid setup - Tutorial 2 style
    nlay, nrow, ncol = 1, 21, 21
    delr = delc = 100.0  # 100m cells
    top, botm = 50.0, 0.0  # 50m thick confined aquifer
    
    # Time setup - 30 days monthly
    nper, perlen, nstp, steady = 1, [30.0], [30], [False]
    
    # Create model
    mf = flopy.modflow.Modflow(modelname='responsive', model_ws='.')
    
    # Discretization
    dis = flopy.modflow.ModflowDis(mf, nlay, nrow, ncol, delr=delr, delc=delc,
                                   top=top, botm=botm, nper=nper,
                                   perlen=perlen, nstp=nstp, steady=steady)
    
    # Boundaries and initial conditions
    ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
    # Constant head boundaries (40m) on all edges
    ibound[0, 0, :] = -1   # North
    ibound[0, -1, :] = -1  # South
    ibound[0, :, 0] = -1   # West
    ibound[0, :, -1] = -1  # East
    
    # Load persistent state or use initial conditions
    state_file = 'responsive_state.pkl'
    if os.path.exists(state_file):
        with open(state_file, 'rb') as f:
            state = pickle.load(f)
            strt = state.get('heads', np.full((nlay, nrow, ncol), 40.0, dtype=np.float32))
            # Ensure boundaries remain at 40m
            strt[0, 0, :] = strt[0, -1, :] = strt[0, :, 0] = strt[0, :, -1] = 40.0
    else:
        strt = np.full((nlay, nrow, ncol), 40.0, dtype=np.float32)
    
    bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)
    
    # Aquifer properties - confined aquifer
    lpf = flopy.modflow.ModflowLpf(mf, hk=hydraulic_k, vka=hydraulic_k, 
                                   sy=0.0, ss=1e-5, laytyp=0, laywet=0)
    
    # Well package
    if abs(well_rate) > 0.001:
        # Ensure well is in interior
        well_row = max(1, min(19, well_row))
        well_col = max(1, min(19, well_col))
        
        # MODFLOW convention: negative = pumping out
        wel_data = [[0, well_row, well_col, -well_rate]]
        wel = flopy.modflow.ModflowWel(mf, stress_period_data={0: wel_data})
    
    # Recharge package
    if abs(recharge_rate) > 0.1:
        cell_area = delr * delc
        active_cells = 19 * 19  # Interior cells only
        recharge_per_cell = recharge_rate / (cell_area * active_cells)
        rch = flopy.modflow.ModflowRch(mf, rech={0: recharge_per_cell})
    
    # Output and solver
    oc = flopy.modflow.ModflowOc(mf, stress_period_data={(0, 29): ['save head']})
    pcg = flopy.modflow.ModflowPcg(mf, mxiter=100, iter1=30, hclose=1e-6)
    
    return mf

def run_modflow_step(*args):
    """
    Main entry point for GoldSim
    
    Args:
        args[0]: well_rate (m³/day, positive = pumping out)
        args[1]: recharge_rate (m³/day total)
        args[2]: well_row (0-20, center is 10)
        args[3]: well_col (0-20, center is 10)
        args[4]: hydraulic_k (m/day)
        args[5]: specific_yield (ignored for confined)
        
    Returns:
        tuple: (northwest_head, center_head, southeast_head) in meters
    """
    global _call_counter, _last_call_time
    
    _call_counter += 1
    current_time = time.time()
    time_gap = current_time - _last_call_time
    
    # Detect new simulation (>10 second gap)
    if time_gap > 10.0 or _call_counter == 1:
        log_message("NEW GOLDSIM SIMULATION DETECTED - Resetting to initial conditions")
        if os.path.exists('responsive_state.pkl'):
            os.remove('responsive_state.pkl')
    
    _last_call_time = current_time
    
    try:
        # Parse inputs with defaults
        well_rate = float(args[0]) if len(args) > 0 else 0.0
        recharge_rate = float(args[1]) if len(args) > 1 else 0.0
        well_row = int(args[2]) if len(args) > 2 else 10
        well_col = int(args[3]) if len(args) > 3 else 10
        hydraulic_k = float(args[4]) if len(args) > 4 else 10.0
        specific_yield = float(args[5]) if len(args) > 5 else 0.18
        
        log_message(f"MODFLOW Call #{_call_counter}: Pump={well_rate:.0f} m³/d at ({well_row},{well_col})")
        
        # Create and run model
        mf = create_modflow_model(well_rate, recharge_rate, well_row, well_col, 
                                 hydraulic_k, specific_yield)
        mf.write_input()
        success, _ = mf.run_model(silent=True)
        
        if not success:
            log_message("MODFLOW failed - using fallback values")
            return (40.0, 40.0, 40.0)
        
        # Read results
        hds = flopy.utils.HeadFile('responsive.hds')
        heads = hds.get_data()
        
        # Extract monitoring heads
        northwest_head = float(heads[0, 4, 6])   # Moved to (4,6)
        center_head = float(heads[0, 10, 10])
        southeast_head = float(heads[0, 13, 13])
        
        # Save state for next month
        state = {'heads': heads, 'timestamp': datetime.now().isoformat()}
        with open('responsive_state.pkl', 'wb') as f:
            pickle.dump(state, f)
        
        log_message(f"Results: NW(4,6)={northwest_head:.2f}m, Center(10,10)={center_head:.2f}m, SE(13,13)={southeast_head:.2f}m")
        
        return (northwest_head, center_head, southeast_head)
        
    except Exception as e:
        log_message(f"ERROR: {str(e)}")
        return (40.0, 40.0, 40.0)

def reset_state():
    """Reset model state for testing"""
    if os.path.exists('responsive_state.pkl'):
        os.remove('responsive_state.pkl')
    print("State reset - ready for fresh simulation")

# Test functionality
if __name__ == "__main__":
    print("Testing MODFLOW Interface...")
    reset_state()
    
    # Test scenarios
    scenarios = [
        (0, 0, 10, 10, 10.0, 0.18, "No pumping"),
        (1000, 0, 10, 10, 10.0, 0.18, "1000 m³/day pumping"),
        (2000, 0, 10, 10, 10.0, 0.18, "2000 m³/day pumping"),
        (1000, 0, 10, 10, 10.0, 0.18, "Back to 1000 m³/day"),
        (0, 0, 10, 10, 10.0, 0.18, "No pumping (recovery)"),
    ]
    
    print("\nPumping | NW Head | Center | SE Head | Description")
    print("(m³/day)|   (m)   |  (m)   |   (m)   |")
    print("-" * 50)
    
    for i, (pump, rech, row, col, k, sy, desc) in enumerate(scenarios):
        if i > 0:
            time.sleep(1)  # Normal monthly interval
            
        result = run_modflow_step(pump, rech, row, col, k, sy)
        nw, center, se = result
        print(f"{pump:7.0f} | {nw:7.2f} | {center:6.2f} | {se:7.2f} | {desc}")
    
    # Test new simulation detection
    print("\nTesting new simulation detection (12 second gap)...")
    time.sleep(12)
    result = run_modflow_step(1500, 0, 10, 10, 10.0, 0.18)
    print(f"   1500 | {result[0]:7.2f} | {result[1]:6.2f} | {result[2]:7.2f} | New simulation (should reset)")
    
    print("\nInterface ready for GoldSim!")