"""
MODFLOW Runner Module - FloPy Tutorial 2 Implementation
Handles MODFLOW execution for FloPy Tutorial 2 benchmark model
"""

import os
import subprocess
import sys
import numpy as np
from state_manager import StateManager
from static_model_setup import create_static_model


class ModflowRunner:
    """FloPy Tutorial 2 MODFLOW simulation runner"""
    
    def __init__(self, gspy_logger):
        self.gspy = gspy_logger
        self.state_manager = StateManager(gspy_logger)
        
    def run_simulation(self, params):
        """
        Run FloPy Tutorial 2 simulation with 1-day time stepping
        
        Args:
            params: Dictionary containing Tutorial 2 parameters:
                - tutorial_period: 1, 2, or 3
                - well_rate_m3_per_day: pumping rate
                - hk: hydraulic conductivity 
                - sy: specific yield
            
        Returns:
            float: Head at center cell (0,4,4) in meters
        """
        try:
            self.gspy.log("Starting FloPy Tutorial 2 simulation...", 2)
            
            # Enhanced parameter logging
            self.gspy.log("DETAILED SIMULATION PARAMETERS:", 2)
            self.gspy.log(f"  - Tutorial Period: {params.get('tutorial_period', 'UNKNOWN')}", 3)
            self.gspy.log(f"  - Well Rate: {params.get('well_rate_m3_per_day', 'UNKNOWN')} m³/day", 3)
            self.gspy.log(f"  - Hydraulic Conductivity: {params.get('hk', 'UNKNOWN')} m/day", 3)
            self.gspy.log(f"  - Specific Yield: {params.get('sy', 'UNKNOWN')}", 3)
            
            # Load previous state if available
            previous_heads = self.state_manager.load_heads()
            if previous_heads is not None:
                self.gspy.log("Loaded persistent heads from previous day", 2)
                # Diagnostic: check head values
                center_head_initial = previous_heads[0, 4, 4]  # [layer, row, col]
                min_head = np.min(previous_heads)
                max_head = np.max(previous_heads)
                avg_head = np.mean(previous_heads)
                self.gspy.log(f"INITIAL HEAD ANALYSIS:", 3)
                self.gspy.log(f"  - Center cell (0,4,4): {center_head_initial:.6f} m", 3)
                self.gspy.log(f"  - Min head in domain: {min_head:.6f} m", 3)
                self.gspy.log(f"  - Max head in domain: {max_head:.6f} m", 3)
                self.gspy.log(f"  - Average head: {avg_head:.6f} m", 3)
            else:
                self.gspy.log("No persistent state found, using initial heads", 2)
            
            # Import FloPy
            try:
                import flopy
                self.gspy.log(f"FloPy version {flopy.__version__} loaded successfully", 2)
            except ImportError as e:
                self.gspy.log(f"FloPy not available: {e}", 0)
                raise
            
            # Get current directory for model files
            model_ws = os.path.dirname(os.path.abspath(__file__))
            
            # Check for MODFLOW executable
            modflow_exe = os.path.join(model_ws, 'mf2005.exe')
            if not os.path.exists(modflow_exe):
                self.gspy.log(f"MODFLOW executable not found at: {modflow_exe}", 1)
                modflow_exe = 'mf2005'  # Try system PATH
            else:
                self.gspy.log(f"MODFLOW executable found: {modflow_exe}", 2)
            
            # Create Tutorial 2 model
            mf = self._create_tutorial_model(params, previous_heads, model_ws, modflow_exe, flopy)
            
            self.gspy.log("Running MODFLOW simulation (1-day step)...", 2)
            
            # Write and run model
            mf.write_input()
            
            # Use hidden subprocess execution
            success, buff = self._run_modflow_hidden(modflow_exe, mf.name, mf.model_ws)
            
            if success:
                self.gspy.log("MODFLOW simulation completed successfully", 2)
                
                # Get center cell head and save state from final time step
                center_head = self._get_center_head(mf)
                # Extract and save heads for persistent state
                self.state_manager.save_heads(mf, 0, 99)  # Period 0, final time step (99)
                
                self.gspy.log(f"Center cell (0,4,4) head: {center_head:.6f} m", 2)
                return center_head
            else:
                self.gspy.log("MODFLOW simulation failed to converge", 0)
                self.gspy.log(f"MODFLOW output: {str(buff)}", 0)
                raise Exception("MODFLOW run failed")
                
        except Exception as e:
            self.gspy.log(f"Simulation error: {str(e)}", 0)
            raise Exception("MODFLOW simulation error")
    
    def _create_tutorial_model(self, params, previous_heads, model_ws, modflow_exe, flopy):
        """
        Create FloPy Tutorial 2 model with GHB and WEL packages
        
        Args:
            params: Dictionary with Tutorial 2 parameters
            previous_heads: Previous head distribution or None
            model_ws: Model workspace directory
            modflow_exe: Path to MODFLOW executable
            flopy: FloPy module
            
        Returns:
            flopy.modflow.Modflow: Tutorial 2 MODFLOW model
        """
        # Extract Tutorial 2 parameters
        tutorial_period = params['tutorial_period']
        well_rate_m3_per_day = params['well_rate_m3_per_day']
        hk = params['hk']
        sy = params['sy']
        
        self.gspy.log(f"Creating Tutorial 2 model for period {tutorial_period}", 2)
        
        # Start with static model (1x10x10 grid, 1000m spacing)
        mf = create_static_model('gspy_tutorial2', model_ws)
        mf.exe_name = modflow_exe
        
        # CRITICAL: Ensure DIS package is configured for DAILY RUNS
        # Must be: nper=1, perlen=[1.0], nstp=[1] for single 1-day step
        mf.remove_package('DIS')
        nlay, nrow, ncol = 1, 10, 10
        delr = delc = 1000.0
        top = 25.0
        botm = 0.0
        
        # MANDATE: Daily time discretization with high resolution for sharp transient capture
        nper = 1           # Single stress period per call
        perlen = [1.0]     # 1 day duration  
        nstp = [100]       # 100 time steps per day for accurate non-linear transient behavior
        steady = [False]   # Transient
        
        dis = flopy.modflow.ModflowDis(mf, nlay=nlay, nrow=nrow, ncol=ncol,
                                      delr=delr, delc=delc, top=top, botm=botm,
                                      nper=nper, perlen=perlen, nstp=nstp, 
                                      steady=steady)
        self.gspy.log(f"DIS configured for DAILY run: nper={nper}, perlen={perlen}, nstp={nstp}", 2)
        
        # Update LPF package with input parameters
        mf.remove_package('LPF')
        ss = 1e-4  # Specific storage
        # CRITICAL: Unconfined aquifer settings for Tutorial 2
        laytyp = 1  # Unconfined aquifer (CRITICAL for transient behavior)
        lpf = flopy.modflow.ModflowLpf(mf, hk=hk, vka=hk, sy=sy, ss=ss, 
                                       laytyp=laytyp, laywet=0)
        self.gspy.log(f"LPF package: K={hk:.1f} m/day, Sy={sy:.3f}", 2)
        self.gspy.log(f"LAYTYP VERIFICATION: {laytyp} (1=unconfined, 0=confined)", 3)
        self.gspy.log(f"  - UNCONFINED AQUIFER: Specific Yield controls transient storage", 3)
        self.gspy.log(f"  - This enables non-linear drawdown response", 3)
        
        # CRITICAL: Update initial heads from persistent state
        self.gspy.log("TUTORIAL 2 BASIC PACKAGE SETUP:", 2)
        mf.remove_package('BAS6')
        ibound = np.ones((1, 10, 10), dtype=np.int32)
        
        # CRITICAL FIX: Ensure GHB boundary cells are NOT set to constant head in IBOUND
        # GHB cells should remain active (ibound=1) so GHB can control them
        self.gspy.log("IBOUND ARRAY ANALYSIS:", 3)
        self.gspy.log(f"  - IBOUND shape: {ibound.shape}", 3)
        self.gspy.log(f"  - All cells set to: {np.unique(ibound)}", 3)
        self.gspy.log(f"  - Center cell (0,4,4) IBOUND: {ibound[0,4,4]}", 3)
        self.gspy.log(f"  - Left boundary cells IBOUND: {ibound[0,:,0]}", 3)
        self.gspy.log(f"  - Right boundary cells IBOUND: {ibound[0,:,9]}", 3)
        
        if previous_heads is not None:
            # Use persistent heads from previous day as initial conditions
            strt = previous_heads
            self.gspy.log("Using persistent heads from previous day as initial conditions", 2)
        else:
            # First run: use default 10.0m initial heads
            strt = np.ones((1, 10, 10), dtype=np.float32) * 10.0
            self.gspy.log("No persistent state found - using default 10.0m initial heads", 2)
            
        # Diagnostic: Check starting heads
        self.gspy.log("STARTING HEAD ANALYSIS:", 3)
        self.gspy.log(f"  - Starting head shape: {strt.shape}", 3)
        self.gspy.log(f"  - Center cell (0,4,4) starting head: {strt[0,4,4]:.6f} m", 3)
        self.gspy.log(f"  - Min starting head: {np.min(strt):.6f} m", 3)
        self.gspy.log(f"  - Max starting head: {np.max(strt):.6f} m", 3)
        
        bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)
        self.gspy.log("BAS package created with active cells throughout domain", 2)
        
        # CRITICAL: Update OC package for daily output control
        mf.remove_package('OC')
        # Save heads at end of 1-day period (final time step) for state persistence
        spd = {(0, 99): ['print head', 'save head']}  # Period 0, final time step 99
        oc = flopy.modflow.ModflowOc(mf, stress_period_data=spd, compact=True)
        self.gspy.log("OC configured to save heads at final time step for persistence", 2)
        
        # Add GHB package based on tutorial_period
        self.gspy.log("TUTORIAL 2 BOUNDARY CONDITIONS SETUP:", 2)
        ghb_cells = []
        if tutorial_period == 1:
            # Period 1: Left and right boundaries both at 10.0m
            for row in range(10):
                ghb_cells.append([0, row, 0, 10.0, 1000000.0])    # Left boundary - VERY HIGH conductance
                ghb_cells.append([0, row, 9, 10.0, 1000000.0])    # Right boundary - VERY HIGH conductance
            self.gspy.log("GHB Period 1: Left=10.0m, Right=10.0m", 2)
            self.gspy.log(f"  - Total GHB cells: {len(ghb_cells)}", 3)
            self.gspy.log(f"  - Conductance: 1000000.0 m²/day (HIGH for fixed-head behavior)", 3)
            
        elif tutorial_period == 2 or tutorial_period == 3:
            # Periods 2&3: Left=10.0m, Right=0.0m
            for row in range(10):
                ghb_cells.append([0, row, 0, 10.0, 1000000.0])    # Left boundary - VERY HIGH conductance
                ghb_cells.append([0, row, 9, 0.0, 1000000.0])     # Right boundary - VERY HIGH conductance
            self.gspy.log("GHB Periods 2&3: Left=10.0m, Right=0.0m", 2)
            self.gspy.log(f"  - Total GHB cells: {len(ghb_cells)}", 3)
            self.gspy.log(f"  - Left boundary head: 10.0m", 3)
            self.gspy.log(f"  - Right boundary head: 0.0m", 3)
            self.gspy.log(f"  - Conductance: 1000000.0 m²/day (HIGH for fixed-head behavior)", 3)
        
        # CRITICAL DIAGNOSTIC: Verify GHB cells don't conflict with IBOUND
        self.gspy.log("GHB-IBOUND COMPATIBILITY CHECK:", 3)
        self.gspy.log(f"  - GHB cells are at boundaries (columns 0 and 9)", 3)
        self.gspy.log(f"  - Center cell (0,4,4) should NOT be in GHB list", 3)
        self.gspy.log(f"  - Interior cells (columns 1-8) should be free to respond", 3)
        
        # Verify center cell is not in GHB list
        center_in_ghb = any(cell[1] == 4 and cell[2] == 4 for cell in ghb_cells)
        self.gspy.log(f"  - Center cell in GHB list: {center_in_ghb} (should be False)", 3)
        
        # Create GHB package
        ghb_spd = {0: ghb_cells}
        ghb = flopy.modflow.ModflowGhb(mf, stress_period_data=ghb_spd)
        self.gspy.log("GHB package created successfully", 2)
        self.gspy.log(f"  - Boundary cells will be controlled by GHB", 3)
        self.gspy.log(f"  - Interior cells remain free to respond to gradients", 3)
        
        # Add WEL package if needed (Period 3 only)
        self.gspy.log("TUTORIAL 2 WELL CONDITIONS SETUP:", 2)
        if tutorial_period == 3 and well_rate_m3_per_day != 0.0:
            # Well at center cell (0,4,4)
            wel_cells = [[0, 4, 4, well_rate_m3_per_day]]
            wel_spd = {0: wel_cells}
            wel = flopy.modflow.ModflowWel(mf, stress_period_data=wel_spd)
            self.gspy.log(f"WEL package: center well = {well_rate_m3_per_day:.1f} m³/day", 2)
            self.gspy.log(f"  - Well location: Layer 0, Row 4, Col 4 (center cell)", 3)
            self.gspy.log(f"  - Pumping rate: {well_rate_m3_per_day:.1f} m³/day", 3)
            self.gspy.log(f"  - Expected behavior: Drawdown at center", 3)
        else:
            self.gspy.log("No well pumping for this period", 2)
            self.gspy.log(f"  - Period: {tutorial_period}", 3)
            self.gspy.log(f"  - Well rate: {well_rate_m3_per_day:.1f} m³/day", 3)
        
        return mf
    
    def _run_modflow_hidden(self, exe_path, model_name, model_ws):
        """
        Run MODFLOW with completely hidden terminal window to prevent flashing
        
        Args:
            exe_path: Path to MODFLOW executable
            model_name: Model name
            model_ws: Model workspace directory
            
        Returns:
            tuple: (success_flag, output_text)
        """
        try:
            # Change to model directory
            old_dir = os.getcwd()
            os.chdir(model_ws)
            
            # Prepare command
            cmd = [exe_path, model_name + '.nam']
            
            # Create startup info for Windows to hide console
            if sys.platform.startswith('win'):
                startupinfo = subprocess.STARTUPINFO()
                startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = subprocess.SW_HIDE
            else:
                startupinfo = None
            
            # Run MODFLOW with hidden console using CREATE_NO_WINDOW flag
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                startupinfo=startupinfo,
                creationflags=subprocess.CREATE_NO_WINDOW if sys.platform.startswith('win') else 0
            )
            
            # Restore directory
            os.chdir(old_dir)
            
            # Check if successful
            success = result.returncode == 0
            output = result.stdout + result.stderr
            
            return success, output
            
        except Exception as e:
            # Restore directory if error occurred
            try:
                os.chdir(old_dir)
            except:
                pass
            return False, str(e)
    
    def _get_center_head(self, mf):
        """
        Get head value at center cell (0,4,4) for Tutorial 2 benchmark
        
        Args:
            mf: MODFLOW model object
            
        Returns:
            float: Head at center cell in meters
        """
        try:
            import flopy
            
            # Diagnostic: Check head file
            head_file_path = os.path.join(mf.model_ws, f'{mf.name}.hds')
            self.gspy.log(f"HEAD FILE ANALYSIS:", 2)
            self.gspy.log(f"  - Path: {head_file_path}", 3)
            self.gspy.log(f"  - Exists: {os.path.exists(head_file_path)}", 3)
            if os.path.exists(head_file_path):
                file_size = os.path.getsize(head_file_path)
                self.gspy.log(f"  - File size: {file_size} bytes", 3)
            
            # Read head file from final time step
            headobj = flopy.utils.HeadFile(head_file_path)
            
            # Check available time steps
            times = headobj.get_times()
            kstpkper = headobj.get_kstpkper()
            self.gspy.log(f"HEAD FILE CONTENTS:", 2)
            self.gspy.log(f"  - Available times: {times}", 3)
            self.gspy.log(f"  - Available kstpkper: {kstpkper}", 3)
            self.gspy.log(f"  - Reading final time step: (99, 0)", 3)
            
            heads = headobj.get_data(kstpkper=(99, 0))  # Final time step (99), period 0
            
            # Comprehensive head analysis
            self.gspy.log(f"FINAL HEAD DISTRIBUTION ANALYSIS:", 2)
            self.gspy.log(f"  - Grid shape: {heads.shape}", 3)
            min_head = np.min(heads)
            max_head = np.max(heads)
            avg_head = np.mean(heads)
            self.gspy.log(f"  - Min head: {min_head:.6f} m", 3)
            self.gspy.log(f"  - Max head: {max_head:.6f} m", 3)
            self.gspy.log(f"  - Average head: {avg_head:.6f} m", 3)
            
            # Get center cell head (layer=0, row=4, col=4)
            center_head = heads[0, 4, 4]
            
            # Analyze surrounding cells for gradient
            self.gspy.log(f"CENTER CELL NEIGHBORHOOD ANALYSIS:", 2)
            for dr in [-1, 0, 1]:
                for dc in [-1, 0, 1]:
                    r, c = 4 + dr, 4 + dc
                    if 0 <= r < heads.shape[1] and 0 <= c < heads.shape[2]:
                        cell_head = heads[0, r, c]
                        if dr == 0 and dc == 0:
                            self.gspy.log(f"  - CENTER (0,{r},{c}): {cell_head:.6f} m <<<", 3)
                        else:
                            self.gspy.log(f"  - Cell (0,{r},{c}): {cell_head:.6f} m", 3)
            
            # Check boundary cells for Tutorial 2 verification
            left_boundary = heads[0, 4, 0]  # Left edge, center row
            right_boundary = heads[0, 4, 9]  # Right edge, center row
            self.gspy.log(f"BOUNDARY CONDITIONS VERIFICATION:", 2)
            self.gspy.log(f"  - Left boundary head: {left_boundary:.6f} m", 3)
            self.gspy.log(f"  - Right boundary head: {right_boundary:.6f} m", 3)
            self.gspy.log(f"  - Head gradient across domain: {right_boundary - left_boundary:.6f} m", 3)
            
            self.gspy.log(f"CENTER CELL RESULT: {center_head:.6f} m", 2)
            return float(center_head)
            
        except Exception as e:
            self.gspy.log(f"Error reading center head: {str(e)}", 0)
            self.gspy.log("Returning fallback head: 10.0 m", 1)
            return 10.0  # Default fallback head