"""
State Manager Module
Handles persistent state management for MODFLOW heads using pickle files
"""

import os
import pickle
import numpy as np


class StateManager:
    """Manages persistent groundwater state between MODFLOW runs"""
    
    def __init__(self, gspy_logger):
        self.gspy = gspy_logger
        self._state_file = self._get_state_file_path()
        
    def _get_state_file_path(self):
        """Get the path to the state file"""
        model_ws = os.path.dirname(os.path.abspath(__file__))
        return os.path.join(model_ws, 'modflow_state.pkl')
    
    def load_heads(self):
        """
        Load previous groundwater heads from state file
        
        Returns:
            numpy.ndarray or None: Previous head distribution or None if not available
        """
        try:
            if os.path.exists(self._state_file):
                with open(self._state_file, 'rb') as f:
                    state = pickle.load(f)
                    heads = state.get('heads', None)
                    if heads is not None:
                        self.gspy.log(f"Loaded persistent heads from: {self._state_file}", 2)
                        return heads
                    else:
                        self.gspy.log("State file exists but contains no head data", 1)
                        return self._generate_default_heads()
            else:
                self.gspy.log("No persistent state file found", 2)
                return self._generate_default_heads()
        except Exception as e:
            self.gspy.log(f"Could not load persistent state: {str(e)}", 1)
            return self._generate_default_heads()
    
    def _generate_default_heads(self):
        """
        Generate FloPy Tutorial 2 initial heads - uniform 10.0m everywhere
        
        Returns:
            numpy.ndarray: Uniform 10.0m head distribution as required by Tutorial 2
        """
        self.gspy.log("Generating FloPy Tutorial 2 initial heads (uniform 10.0m)", 2)
        
        # Grid dimensions as specified in requirements (10x10)
        nlay, nrow, ncol = 1, 10, 10
        
        # FloPy Tutorial 2 REQUIREMENT: Uniform 10.0m initial heads everywhere
        ihead = np.full((nlay, nrow, ncol), 10.0, dtype=np.float32)
        
        self.gspy.log("Tutorial 2 initial heads created: uniform 10.0m everywhere", 2)
        return ihead
    
    def save_heads(self, mf, period, timestep):
        """
        Save current groundwater heads to state file
        
        Args:
            mf: MODFLOW model object
            period: Stress period
            timestep: Time step
        """
        try:
            import flopy
            self.gspy.log("Extracting heads for persistent state...", 2)
            
            # Read final heads from MODFLOW output
            headobj = flopy.utils.HeadFile(os.path.join(mf.model_ws, f'{mf.name}.hds'))
            heads = headobj.get_data(kstpkper=(timestep, period))
            
            # Save to pickle file
            state = {'heads': heads}
            with open(self._state_file, 'wb') as f:
                pickle.dump(state, f)
            
            self.gspy.log(f"Groundwater heads saved to: {self._state_file}", 2)
            
        except Exception as e:
            self.gspy.log(f"Could not save persistent state: {str(e)}", 1)
    
    def reset_state(self):
        """Reset persistent state - useful for testing"""
        try:
            if os.path.exists(self._state_file):
                os.remove(self._state_file)
                self.gspy.log(f"Persistent state file removed: {self._state_file}", 2)
            else:
                self.gspy.log("No persistent state file to remove", 2)
        except Exception as e:
            self.gspy.log(f"Could not remove persistent state file: {str(e)}", 1)