"""
Monthly MODFLOW-GoldSim Coupling Prototype
==========================================
Demonstrates integrated surface-groundwater management with monthly decision cycles.

Story: Water utility makes monthly decisions about:
- Surface water diversions from river
- Groundwater pumping rates  
- Environmental flow requirements

MODFLOW runs 30-day simulations and reports:
- Average monthly return flows to river
- End-of-month groundwater levels
- Monthly water budget components
"""

import os
import numpy as np
import flopy
import pickle
from datetime import datetime, timedelta
import json

class MonthlyGroundwaterManager:
    """Monthly groundwater management with MODFLOW coupling"""
    
    def __init__(self, model_ws='.'):
        self.model_ws = model_ws
        self.state_file = os.path.join(model_ws, 'monthly_state.pkl')
        self.results_file = os.path.join(model_ws, 'monthly_results.json')
        self.current_month = 1
        self.simulation_year = 2024
        
        # Initialize results tracking
        self.monthly_results = []
        
    def create_monthly_model(self):
        """Create MODFLOW model for monthly 30-day simulations"""
        
        # Model dimensions - realistic water utility scale
        nlay, nrow, ncol = 1, 25, 25  # 25x25 grid
        delr = delc = 200.0  # 200m cells = 5km x 5km domain
        top = 50.0
        botm = 0.0
        
        # Monthly temporal discretization
        nper = 1
        perlen = [30.0]  # 30 days per stress period
        nstp = [30]      # Daily time steps within month
        steady = [False]
        
        # Create model
        mf = flopy.modflow.Modflow(modelname='monthly_gw', model_ws=self.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 - active domain with constant head boundaries
        ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
        # Set boundary conditions: constant head on edges
        ibound[0, 0, :] = -1   # North boundary
        ibound[0, -1, :] = -1  # South boundary  
        ibound[0, :, 0] = -1   # West boundary
        ibound[0, :, -1] = -1  # East boundary
        
        # Load persistent heads or use initial conditions
        if os.path.exists(self.state_file):
            with open(self.state_file, 'rb') as f:
                state = pickle.load(f)
                strt = state.get('heads', self._get_initial_heads())
        else:
            strt = self._get_initial_heads()
            
        bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)
        
        # LPF package - unconfined aquifer
        hk = 25.0     # Higher K for faster response
        vka = 25.0
        sy = 0.15     # Typical unconfined aquifer
        ss = 1e-4
        
        lpf = flopy.modflow.ModflowLpf(mf, hk=hk, vka=vka, sy=sy, ss=ss,
                                       laytyp=1, laywet=0)
        
        # OC package - save heads and budget
        spd = {}
        for kper in range(nper):
            for kstp in range(nstp[kper]):
                spd[(kper, kstp)] = ['save head', 'save budget']
        
        oc = flopy.modflow.ModflowOc(mf, stress_period_data=spd, compact=True)
        
        # PCG solver
        pcg = flopy.modflow.ModflowPcg(mf, mxiter=500, iter1=50)
        
        return mf
    
    def _get_initial_heads(self):
        """Create realistic initial head distribution"""
        nlay, nrow, ncol = 1, 25, 25
        
        # Create gentle gradient from NE to SW (typical regional flow)
        heads = np.zeros((nlay, nrow, ncol))
        
        for i in range(nrow):
            for j in range(ncol):
                # Gradient from 45m (NE) to 35m (SW)
                heads[0, i, j] = 45.0 - 0.4 * i - 0.4 * j
                
        # Add local high around pumping well location
        center_i, center_j = 12, 12
        for i in range(nrow):
            for j in range(ncol):
                dist = np.sqrt((i - center_i)**2 + (j - center_j)**2)
                if dist < 5:
                    heads[0, i, j] += 2.0 * np.exp(-dist/3.0)
        
        return heads.astype(np.float32)
    
    def run_monthly_simulation(self, surface_diversion_m3d, pumping_rate_m3d, 
                              recharge_rate_md):
        """
        Run monthly MODFLOW simulation
        
        Args:
            surface_diversion_m3d: Surface water diversion (m³/day)
            pumping_rate_m3d: Groundwater pumping rate (m³/day, positive = pumping)
            recharge_rate_md: Distributed recharge rate (m/day)
            
        Returns:
            dict: Monthly results including return flows, heads, water budget
        """
        
        # Create model
        mf = self.create_monthly_model()
        
        # Add recharge package
        if recharge_rate_md > 0:
            rch_data = {0: recharge_rate_md}  # Apply to stress period 0
            rch = flopy.modflow.ModflowRch(mf, rech=rch_data)
        
        # Add river package (represents managed river system)
        river_cells = self._get_river_cells()
        if len(river_cells) > 0:
            # River stage varies with surface diversion
            base_stage = 38.0
            # More diversion = lower river stage
            river_stage = base_stage - (surface_diversion_m3d / 10000.0)
            
            riv_data = []
            for i, j in river_cells:
                # High conductance for good connection
                conductance = 1000.0  # m²/day
                rbot = 35.0  # River bottom
                riv_data.append([0, i, j, river_stage, conductance, rbot])
            
            riv_spd = {0: riv_data}
            riv = flopy.modflow.ModflowRiv(mf, stress_period_data=riv_spd)
        
        # Add well package (municipal pumping)
        if pumping_rate_m3d > 0:
            # Central well location
            well_i, well_j = 12, 12
            wel_data = [[0, well_i, well_j, -pumping_rate_m3d]]  # Negative = pumping
            wel_spd = {0: wel_data}
            wel = flopy.modflow.ModflowWel(mf, stress_period_data=wel_spd)
        
        # Write and run model
        mf.write_input()
        success, buff = mf.run_model(silent=True)
        
        if not success:
            raise RuntimeError("MODFLOW simulation failed")
        
        # Process results
        results = self._process_monthly_results(mf, surface_diversion_m3d, 
                                               pumping_rate_m3d, recharge_rate_md)
        
        # Save state for next month
        self._save_monthly_state(mf)
        
        # Track results
        self.monthly_results.append(results)
        self._save_results_log()
        
        return results
    
    def _get_river_cells(self):
        """Define river cells - diagonal across domain"""
        river_cells = []
        # Diagonal river from NW to SE
        for i in range(5, 20):
            j = i
            if j < 25:
                river_cells.append((i, j))
        return river_cells
    
    def _process_monthly_results(self, mf, surface_diversion, pumping_rate, recharge_rate):
        """Process MODFLOW results into monthly summary"""
        
        # Read heads
        hds_file = os.path.join(self.model_ws, f'{mf.name}.hds')
        hds = flopy.utils.HeadFile(hds_file)
        heads = hds.get_data()  # Get final heads
        
        # Read budget
        cbc_file = os.path.join(self.model_ws, f'{mf.name}.cbc')
        cbc = flopy.utils.CellBudgetFile(cbc_file)
        
        # Calculate return flows (river leakage)
        river_flow = 0.0
        try:
            riv_data = cbc.get_data(text='RIVER LEAKAGE')[0]
            river_flow = np.sum(riv_data)  # Positive = discharge to river
        except:
            pass
        
        # Calculate water budget components
        storage_change = 0.0
        try:
            sto_data = cbc.get_data(text='STORAGE')[0]
            storage_change = np.sum(sto_data)
        except:
            pass
        
        # Monthly summary
        month_name = self._get_month_name(self.current_month)
        
        results = {
            'month': self.current_month,
            'month_name': month_name,
            'year': self.simulation_year,
            'inputs': {
                'surface_diversion_m3d': surface_diversion,
                'pumping_rate_m3d': pumping_rate,
                'recharge_rate_md': recharge_rate
            },
            'outputs': {
                'avg_return_flow_m3d': float(river_flow),
                'avg_return_flow_m3s': float(river_flow / 86400.0),
                'end_month_head_center_m': float(heads[0, 12, 12]),
                'end_month_head_avg_m': float(np.mean(heads)),
                'storage_change_m3': float(storage_change),
                'water_balance_error_pct': self._calculate_water_balance_error(
                    recharge_rate, pumping_rate, river_flow, storage_change)
            },
            'timestamp': datetime.now().isoformat()
        }
        
        self.current_month += 1
        if self.current_month > 12:
            self.current_month = 1
            self.simulation_year += 1
            
        return results
    
    def _calculate_water_balance_error(self, recharge_rate, pumping_rate, 
                                     river_flow, storage_change):
        """Calculate water balance error as percentage"""
        domain_area = (25 * 200) * (25 * 200)  # m²
        
        inflow = recharge_rate * domain_area * 30  # m³/month
        outflow = pumping_rate * 30 + river_flow * 30  # m³/month
        
        if abs(inflow + outflow) > 1e-6:
            error = abs((inflow - outflow - storage_change) / (inflow + outflow)) * 100
        else:
            error = 0.0
            
        return float(error)
    
    def _get_month_name(self, month_num):
        """Convert month number to name"""
        months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun',
                 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
        return months[month_num - 1]
    
    def _save_monthly_state(self, mf):
        """Save groundwater state for next month"""
        hds_file = os.path.join(self.model_ws, f'{mf.name}.hds')
        hds = flopy.utils.HeadFile(hds_file)
        final_heads = hds.get_data()
        
        state = {
            'heads': final_heads,
            'month': self.current_month,
            'year': self.simulation_year,
            'timestamp': datetime.now().isoformat()
        }
        
        with open(self.state_file, 'wb') as f:
            pickle.dump(state, f)
    
    def _save_results_log(self):
        """Save results log for analysis"""
        with open(self.results_file, 'w') as f:
            json.dump(self.monthly_results, f, indent=2)
    
    def reset_simulation(self):
        """Reset to start new annual simulation"""
        if os.path.exists(self.state_file):
            os.remove(self.state_file)
        self.current_month = 1
        self.simulation_year = 2024
        self.monthly_results = []
        print("Monthly simulation reset - ready for new year")


# GSPy Interface for Monthly Coupling
def run_monthly_modflow(*args):
    """
    GSPy entry point for monthly coupling
    
    Args:
        args[0]: surface_diversion_m3d (float) - Surface water diversion rate
        args[1]: pumping_rate_m3d (float) - Groundwater pumping rate  
        args[2]: recharge_rate_md (float) - Distributed recharge rate
        
    Returns:
        tuple: (return_flow_m3s, center_head_m, avg_head_m)
    """
    try:
        surface_diversion = float(args[0])
        pumping_rate = float(args[1]) 
        recharge_rate = float(args[2])
        
        # Create manager and run simulation
        manager = MonthlyGroundwaterManager()
        results = manager.run_monthly_simulation(surface_diversion, pumping_rate, recharge_rate)
        
        # Return key outputs for GoldSim
        return_flow_m3s = results['outputs']['avg_return_flow_m3s']
        center_head = results['outputs']['end_month_head_center_m']
        avg_head = results['outputs']['end_month_head_avg_m']
        
        return (return_flow_m3s, center_head, avg_head)
        
    except Exception as e:
        print(f"Error in monthly MODFLOW: {e}")
        return (0.0, 40.0, 40.0)  # Safe fallback values


if __name__ == "__main__":
    print("Testing Monthly MODFLOW-GoldSim Coupling...")
    
    # Test annual cycle
    manager = MonthlyGroundwaterManager()
    manager.reset_simulation()
    
    # Simulate seasonal patterns
    seasonal_data = [
        # Month, Surface Diversion, Pumping, Recharge
        (1, 5000, 2000, 0.001),   # Jan: Low demand, winter recharge
        (2, 4000, 1800, 0.002),   # Feb: Low demand, winter recharge  
        (3, 6000, 2200, 0.003),   # Mar: Spring increase
        (4, 8000, 2800, 0.002),   # Apr: Spring demand
        (5, 12000, 3500, 0.001),  # May: High demand starts
        (6, 15000, 4000, 0.0005), # Jun: Peak demand, low recharge
    ]
    
    print("\nRunning 6-month test simulation...")
    for month, div, pump, rech in seasonal_data:
        results = manager.run_monthly_simulation(div, pump, rech)
        print(f"{results['month_name']}: Return flow = {results['outputs']['avg_return_flow_m3s']:.3f} m³/s, "
              f"Head = {results['outputs']['end_month_head_center_m']:.1f} m")
    
    print(f"\nResults saved to: {manager.results_file}")