xinanjiang_model.py

"""
Xinanjiang (新安江) Hydrological Model

The Xinanjiang model is a conceptual rainfall-runoff model developed by Zhao Ren-jun
in 1973 at Hohai University, China. It is widely used in humid and semi-humid regions
for flood forecasting and water resources management.

Mathematical Foundation:
=======================

1. Evapotranspiration Calculation:
   - Three-layer evapotranspiration structure (upper, lower, and deep layer)
   - EU = EP * (W / WM)                    if W < WM
   - EL = EP * (W - WM) / (WM * (C - 1))  if W >= WM
   - ED = EP * (W - WM - WLM) / (WM * (C - 1) * (C - 1)) if W >= WM + WLM

2. Runoff Generation (Saturation Excess Mechanism):
   The model assumes a parabolic distribution of soil moisture capacity:
   
   f = SM * (1 - (1 - A)^(1/(1+B)))
   
   Where:
   - f: cumulative frequency distribution of soil moisture capacity
   - SM: areal mean soil moisture storage capacity (mm)
   - A: area fraction with soil moisture capacity <= a
   - B: shape parameter of the distribution curve

   Runoff generation:
   R = P + A0 - SM + SM * (1 - (P + A0)/SM)^(1+B)  if P + A0 < SM
   R = P + A0 - SM                                   if P + A0 >= SM

3. Runoff Separation:
   Total runoff is separated into surface runoff (RS), interflow (RI), and groundwater (RG)
   
   RS = PE * FR                    Surface runoff
   RI = PE * (1 - FR) * KI        Interflow
   RG = PE * (1 - FR) * (1 - KI)  Groundwater

4. Flow Concentration:
   - Surface runoff uses unit hydrograph
   - Interflow and groundwater use linear reservoirs
   
   Q(t) = KS * S(t)   for surface
   Q(t) = KI * RI(t)  for interflow  
   Q(t) = KG * RG(t)  for groundwater

Author: Zhao Ren-jun (Original), Implementation by [Your Name]
Date: 2024
"""

import numpy as np
from typing import Tuple, Dict
import warnings
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime, timedelta
import os
import seaborn as sns
from scipy.stats import pearsonr
from scipy import stats

# Configure matplotlib for Chinese font display / 配置matplotlib以显示中文
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'STSong', 'KaiTi', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False  # Fix minus sign display / 修复负号显示


class XinanjiangModel:
    """
    Xinanjiang (新安江) rainfall-runoff model implementation.
    
    This model simulates the rainfall-runoff process using a saturation excess 
    runoff generation mechanism, which is particularly suitable for humid regions.
    
    Parameters:
    -----------
    K : float
        Ratio of potential evapotranspiration to pan evaporation (typically 0.7-1.2)
    B : float
        Exponent of the tension water capacity curve (typically 0.1-0.4)
    IMP : float
        Impervious area fraction (0-0.1)
    WM : float
        Average soil moisture storage capacity (mm, typically 120-200)
    WUM : float
        Upper layer soil moisture storage capacity (mm, typically 20-50)
    WLM : float
        Lower layer soil moisture storage capacity (mm, typically 60-90)
    C : float
        Coefficient of deep layer evapotranspiration (typically 0.15-0.20)
    SM : float
        Areal mean free water capacity of surface soil layer (mm, typically 10-50)
    EX : float
        Exponent of the free water capacity curve (typically 1.0-1.5)
    KI : float
        Outflow coefficient of interflow storage (typically 0.2-0.7)
    KG : float
        Outflow coefficient of groundwater storage (typically 0.2-0.7)
    CI : float
        Recession constant of interflow storage (typically 0.5-0.9)
    CG : float
        Recession constant of groundwater storage (typically 0.95-0.998)
    """
    
    def __init__(self, 
                 K: float = 1.0,
                 B: float = 0.3,
                 IMP: float = 0.01,
                 WM: float = 150.0,
                 WUM: float = 30.0,
                 WLM: float = 70.0,
                 C: float = 0.17,
                 SM: float = 30.0,
                 EX: float = 1.5,
                 KI: float = 0.3,
                 KG: float = 0.3,
                 CI: float = 0.7,
                 CG: float = 0.98):
        
        # Model parameters
        self.K = K      # Evapotranspiration coefficient
        self.B = B      # Exponent of tension water distribution curve
        self.IMP = IMP  # Impervious area ratio
        self.WM = WM    # Average tension water capacity
        self.WUM = WUM  # Upper layer tension water capacity
        self.WLM = WLM  # Lower layer tension water capacity
        self.WDM = WM - WUM - WLM  # Deep layer capacity
        self.C = C      # Deep layer evapotranspiration coefficient
        self.SM = SM    # Areal mean free water capacity
        self.EX = EX    # Exponent of free water distribution curve
        self.KI = KI    # Interflow outflow coefficient
        self.KG = KG    # Groundwater outflow coefficient
        self.CI = CI    # Interflow recession constant
        self.CG = CG    # Groundwater recession constant
        
        # State variables (initial conditions)
        self.W = WM * 0.6    # Initial soil moisture (60% of capacity)
        self.S = SM * 0.3    # Initial free water storage
        self.SI = 0.0        # Initial interflow storage
        self.SG = 0.0        # Initial groundwater storage
        
    def evapotranspiration(self, EP: float) -> Tuple[float, float, float]:
        """
        Calculate three-layer evapotranspiration.
        
        EP: Potential evapotranspiration (mm)
        Returns: (EU, EL, ED) - Upper, Lower, Deep layer evapotranspiration
        """
        EP = self.K * EP  # Adjust potential ET
        
        EU = EL = ED = 0.0
        
        # Upper layer evapotranspiration
        if self.W < self.WUM:
            EU = min(EP * self.W / self.WUM, self.W)
        else:
            EU = EP
            
        # Lower layer evapotranspiration
        if self.W >= self.WUM:
            W_temp = self.W - self.WUM
            if W_temp < self.WLM:
                EL = min((EP - EU) * W_temp / self.WLM, W_temp)
            else:
                EL = EP - EU
                
        # Deep layer evapotranspiration
        if self.W >= (self.WUM + self.WLM):
            W_temp = self.W - self.WUM - self.WLM
            if W_temp < self.WDM:
                ED = min(self.C * (EP - EU - EL) * W_temp / self.WDM, W_temp)
            else:
                ED = self.C * (EP - EU - EL)
        
        return EU, EL, ED
    
    def runoff_generation(self, P: float, EP: float) -> float:
        """
        Calculate runoff generation using saturation excess mechanism.
        
        P: Precipitation (mm)
        EP: Potential evapotranspiration (mm)
        Returns: Runoff (mm)
        """
        # Calculate evapotranspiration
        EU, EL, ED = self.evapotranspiration(EP)
        E = EU + EL + ED
        
        # Update soil moisture
        PE = max(0, P - E)  # Net precipitation
        
        # Calculate runoff from pervious area using parabolic curve
        if PE > 0 and self.WM > 0.01:  # Add safety check for division
            A = self.WM * (1.0 - (1.0 - self.W / self.WM) ** (1.0 / (1.0 + self.B)))
            
            if PE + A < self.WM:
                # Partial area generates runoff
                R = PE + A - self.WM + self.WM * (1.0 - (PE + A) / self.WM) ** (1.0 + self.B)
            else:
                # Entire area generates runoff
                R = PE + A - self.WM
                
            # Update soil moisture
            self.W = min(self.WM, self.W + PE - R)
        else:
            R = 0.0
            self.W = max(0.0, self.W - E)
        
        # Add runoff from impervious area
        R_total = R * (1.0 - self.IMP) + P * self.IMP
        
        return R_total
    
    def runoff_separation(self, R: float) -> Tuple[float, float, float]:
        """
        Separate total runoff into surface runoff, interflow, and groundwater.
        
        R: Total runoff (mm)
        Returns: (RS, RI, RG) - Surface runoff, Interflow, Groundwater
        """
        if R <= 0:
            return 0.0, 0.0, 0.0
        
        # Update free water storage
        self.S = self.S + R
        
        # Calculate surface runoff using free water capacity curve
        if self.S <= self.SM:
            # Free water capacity not exceeded - calculate using distribution
            S_ratio = self.S / self.SM
            # Surface runoff occurs when free water exceeds local capacity
            if S_ratio > 0 and self.EX > 0:
                RS = max(0, self.S - self.SM * (1.0 - S_ratio ** (1.0 + self.EX)))
            else:
                RS = 0.0
        else:
            # Free water capacity exceeded
            RS = self.S - self.SM
        
        # Limit surface runoff to available runoff
        RS = min(RS, R)
        RS = max(0, RS)
        
        # Update storage after surface runoff
        self.S = max(0.0, self.S - RS)
        
        # Remaining water for interflow and groundwater
        RSS = max(0, R - RS)
        
        # Separate interflow and groundwater
        if RSS > 0:
            RI = self.KI * RSS
            RG = (1.0 - self.KI) * RSS
        else:
            RI = 0.0
            RG = 0.0
        
        return RS, RI, RG
    
    def flow_routing(self, RS: float, RI: float, RG: float) -> Tuple[float, float, float]:
        """
        Route surface runoff, interflow, and groundwater to outlet.
        
        Uses linear reservoir routing for interflow and groundwater.
        
        Returns: (QS, QI, QG) - Surface flow, Interflow, Groundwater flow
        """
        # Surface runoff (no routing in this simplified version)
        QS = RS
        
        # Interflow routing (linear reservoir)
        self.SI = self.CI * self.SI + RI
        QI = (1.0 - self.CI) * self.SI
        
        # Groundwater routing (linear reservoir)
        self.SG = self.CG * self.SG + RG
        QG = (1.0 - self.CG) * self.SG
        
        return QS, QI, QG
    
    def run_timestep(self, P: float, EP: float) -> Dict[str, float]:
        """
        Run the model for one time step.
        
        Parameters:
        -----------
        P : float
            Precipitation (mm)
        EP : float
            Potential evapotranspiration (mm)
            
        Returns:
        --------
        dict : Dictionary containing model outputs
            - Q: Total discharge (mm)
            - QS: Surface runoff (mm)
            - QI: Interflow (mm)
            - QG: Groundwater flow (mm)
            - R: Total runoff generated (mm)
            - E: Actual evapotranspiration (mm)
            - W: Soil moisture (mm)
            - S: Free water storage (mm)
        """
        # Calculate evapotranspiration
        EU, EL, ED = self.evapotranspiration(EP)
        E = EU + EL + ED
        
        # Generate runoff
        R = self.runoff_generation(P, EP)
        
        # Separate runoff components
        RS, RI, RG = self.runoff_separation(R)
        
        # Route flows
        QS, QI, QG = self.flow_routing(RS, RI, RG)
        
        # Total discharge
        Q = QS + QI + QG
        
        return {
            'Q': Q,      # Total discharge
            'QS': QS,    # Surface runoff
            'QI': QI,    # Interflow
            'QG': QG,    # Groundwater flow
            'R': R,      # Total runoff generated
            'E': E,      # Evapotranspiration
            'W': self.W, # Soil moisture
            'S': self.S  # Free water storage
        }
    
    def run(self, P: np.ndarray, EP: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Run the model for multiple time steps.
        
        Parameters:
        -----------
        P : np.ndarray
            Precipitation time series (mm)
        EP : np.ndarray
            Potential evapotranspiration time series (mm)
            
        Returns:
        --------
        dict : Dictionary containing model outputs as arrays
        """
        n_steps = len(P)
        
        # Initialize output arrays
        Q = np.zeros(n_steps)
        QS = np.zeros(n_steps)
        QI = np.zeros(n_steps)
        QG = np.zeros(n_steps)
        R = np.zeros(n_steps)
        E = np.zeros(n_steps)
        W = np.zeros(n_steps)
        S = np.zeros(n_steps)
        
        # Run model for each time step
        for t in range(n_steps):
            result = self.run_timestep(P[t], EP[t])
            Q[t] = result['Q']
            QS[t] = result['QS']
            QI[t] = result['QI']
            QG[t] = result['QG']
            R[t] = result['R']
            E[t] = result['E']
            W[t] = result['W']
            S[t] = result['S']
        
        return {
            'Q': Q,
            'QS': QS,
            'QI': QI,
            'QG': QG,
            'R': R,
            'E': E,
            'W': W,
            'S': S
        }


def create_xinanjiang_plots(model, P, EP, results, save_dir="figures"):
    """
    Create comprehensive visualization plots for Xinanjiang model results.
    
    Parameters:
    -----------
    model : XinanjiangModel
        The model instance
    P, EP : np.ndarray
        Input precipitation and evapotranspiration
    results : dict
        Model simulation results
    save_dir : str
        Directory to save figures
    """
    # Ensure save directory exists
    os.makedirs(save_dir, exist_ok=True)
    
    # Set style
    plt.style.use('seaborn-v0_8-darkgrid')
    sns.set_palette("husl")
    
    # Re-configure Chinese font after style setting / 样式设置后重新配置中文字体
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'STSong', 'KaiTi', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    
    # Create date index
    n_days = len(P)
    dates = [datetime(2020, 1, 1) + timedelta(days=i) for i in range(n_days)]
    
    # Figure 1: Comprehensive time series
    fig, axes = plt.subplots(5, 1, figsize=(15, 14), sharex=True)
    fig.suptitle('Xinanjiang (新安江) Model - Comprehensive Analysis', fontsize=16, fontweight='bold')
    
    # Precipitation (inverted)
    axes[0].bar(dates, P, color='steelblue', alpha=0.7, width=1)
    axes[0].set_ylabel('Precipitation\n(mm/day)', fontweight='bold')
    axes[0].invert_yaxis()
    axes[0].grid(True, alpha=0.3)
    axes[0].set_ylim(max(P) * 1.1, 0)
    
    # Evapotranspiration
    axes[1].plot(dates, EP, color='orange', linewidth=1.5, label='Potential ET', alpha=0.8)
    axes[1].plot(dates, results['E'], color='red', linewidth=1.5, label='Actual ET', alpha=0.8)
    axes[1].fill_between(dates, EP, alpha=0.2, color='orange')
    axes[1].fill_between(dates, results['E'], alpha=0.3, color='red')
    axes[1].set_ylabel('Evapotranspiration\n(mm/day)', fontweight='bold')
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    # Discharge components
    axes[2].plot(dates, results['Q'], color='blue', linewidth=2, label='Total Discharge')
    axes[2].plot(dates, results['QS'], color='red', linewidth=1.5, label='Surface Runoff', alpha=0.8)
    axes[2].plot(dates, results['QI'], color='green', linewidth=1.5, label='Interflow', alpha=0.8)
    axes[2].plot(dates, results['QG'], color='purple', linewidth=1.5, label='Groundwater', alpha=0.8)
    axes[2].set_ylabel('Discharge\n(mm/day)', fontweight='bold')
    axes[2].grid(True, alpha=0.3)
    axes[2].legend()
    
    # Soil moisture states
    axes[3].plot(dates, results['W'], color='brown', linewidth=2, label='Total Soil Moisture')
    # Only plot available moisture components
    if 'WU' in results and 'WL' in results and 'WD' in results:
        axes[3].plot(dates, results['WU'], color='orange', linewidth=1.5, label='Upper Layer', alpha=0.8)
        axes[3].plot(dates, results['WL'], color='green', linewidth=1.5, label='Lower Layer', alpha=0.8)
        axes[3].plot(dates, results['WD'], color='blue', linewidth=1.5, label='Deep Layer', alpha=0.8)
    axes[3].axhline(y=model.WM, color='red', linestyle='--', alpha=0.7, label=f'Capacity (WM={model.WM}mm)')
    axes[3].fill_between(dates, results['W'], alpha=0.2, color='brown')
    axes[3].set_ylabel('Soil Moisture\n(mm)', fontweight='bold')
    axes[3].grid(True, alpha=0.3)
    axes[3].legend()
    
    # Runoff generation
    axes[4].plot(dates, results['R'], color='darkblue', linewidth=2, label='Total Runoff')
    # Only plot available runoff generation components
    if 'RS' in results and 'RI' in results and 'RG' in results:
        axes[4].plot(dates, results['RS'], color='red', linewidth=1.5, label='Surface Runoff Gen.', alpha=0.8)
        axes[4].plot(dates, results['RI'], color='green', linewidth=1.5, label='Interflow Gen.', alpha=0.8)
        axes[4].plot(dates, results['RG'], color='purple', linewidth=1.5, label='Groundwater Gen.', alpha=0.8)
    axes[4].fill_between(dates, results['R'], alpha=0.2, color='darkblue')
    axes[4].set_ylabel('Runoff Generation\n(mm/day)', fontweight='bold')
    axes[4].set_xlabel('Date', fontweight='bold')
    axes[4].grid(True, alpha=0.3)
    axes[4].legend()
    
    # Format x-axis
    for ax in axes:
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'xinanjiang_comprehensive.png'), dpi=300, bbox_inches='tight')
    plt.close()
    
    # Figure 2: Water balance and model structure
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle('Xinanjiang Model - Water Balance & Structure Analysis', fontsize=16, fontweight='bold')
    
    # Cumulative water balance
    cum_P = np.cumsum(P)
    cum_E = np.cumsum(results['E'])
    cum_Q = np.cumsum(results['Q'])
    cum_R = np.cumsum(results['R'])
    
    axes[0, 0].plot(dates, cum_P, label='Cumulative Precipitation', linewidth=2, color='blue')
    axes[0, 0].plot(dates, cum_E, label='Cumulative Evapotranspiration', linewidth=2, color='orange')
    axes[0, 0].plot(dates, cum_Q, label='Cumulative Discharge', linewidth=2, color='red')
    axes[0, 0].plot(dates, cum_R, label='Cumulative Runoff', linewidth=2, color='green')
    axes[0, 0].set_title('Cumulative Water Balance', fontweight='bold')
    axes[0, 0].set_ylabel('Cumulative Water (mm)', fontweight='bold')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    
    # Runoff components pie chart
    total_qs = np.sum(results['QS'])
    total_qi = np.sum(results['QI'])
    total_qg = np.sum(results['QG'])
    
    labels = ['Surface Runoff', 'Interflow', 'Groundwater']
    sizes = [total_qs, total_qi, total_qg]
    colors = ['red', 'green', 'purple']
    
    axes[0, 1].pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
    axes[0, 1].set_title('Discharge Components Distribution', fontweight='bold')
    
    # Soil moisture vs precipitation relationship
    axes[1, 0].scatter(P, results['W'], alpha=0.6, c=results['E'], cmap='viridis')
    axes[1, 0].set_xlabel('Precipitation (mm/day)', fontweight='bold')
    axes[1, 0].set_ylabel('Soil Moisture (mm)', fontweight='bold')
    axes[1, 0].set_title('Soil Moisture vs Precipitation', fontweight='bold')
    cbar = plt.colorbar(axes[1, 0].collections[0], ax=axes[1, 0])
    cbar.set_label('Actual ET (mm/day)', fontweight='bold')
    axes[1, 0].grid(True, alpha=0.3)
    
    # ET efficiency (Actual/Potential)
    et_efficiency = results['E'] / (EP + 1e-6)  # Avoid division by zero
    et_efficiency = np.clip(et_efficiency, 0, 1)  # Clip to valid range
    
    axes[1, 1].scatter(results['W']/model.WM, et_efficiency, alpha=0.6, c=P, cmap='Blues')
    axes[1, 1].set_xlabel('Soil Moisture Ratio (W/WM)', fontweight='bold')
    axes[1, 1].set_ylabel('ET Efficiency (E/EP)', fontweight='bold')
    axes[1, 1].set_title('ET Efficiency vs Soil Moisture', fontweight='bold')
    axes[1, 1].axhline(y=1.0, color='red', linestyle='--', alpha=0.7, label='Maximum Efficiency')
    axes[1, 1].legend()
    cbar2 = plt.colorbar(axes[1, 1].collections[0], ax=axes[1, 1])
    cbar2.set_label('Precipitation (mm/day)', fontweight='bold')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'xinanjiang_water_balance.png'), dpi=300, bbox_inches='tight')
    plt.close()
    
    # Figure 3: Model parameters and performance
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle('Xinanjiang Model - Parameters & Performance Analysis', fontsize=16, fontweight='bold')
    
    # Model parameters visualization
    param_names = ['K', 'B', 'IMP×100', 'WM/10', 'WUM', 'WLM', 'C×100', 'SM', 'EX', 'KI×10', 'KG×10']
    param_values = [model.K, model.B, model.IMP*100, model.WM/10, model.WUM, model.WLM, 
                   model.C*100, model.SM, model.EX, model.KI*10, model.KG*10]
    
    # Use modern Matplotlib colormap API to avoid deprecation warnings
    from matplotlib import colormaps as mcm
    bars = axes[0, 0].bar(
        range(len(param_names)),
        param_values,
        color=mcm.get_cmap('Set3')(np.linspace(0, 1, len(param_names)))
    )
    axes[0, 0].set_title('Model Parameters (scaled for visualization)', fontweight='bold')
    axes[0, 0].set_ylabel('Parameter Value', fontweight='bold')
    axes[0, 0].set_xticks(range(len(param_names)))
    axes[0, 0].set_xticklabels(param_names, rotation=45)
    axes[0, 0].grid(True, alpha=0.3)
    
    # Add value labels on bars
    for bar, value in zip(bars, param_values):
        height = bar.get_height()
        axes[0, 0].text(bar.get_x() + bar.get_width()/2., height + height*0.01,
                       f'{value:.2f}', ha='center', va='bottom', fontsize=8)
    
    # Flow duration curve
    q_sorted = np.sort(results['Q'])[::-1]
    exceedance = np.arange(1, len(q_sorted) + 1) / len(q_sorted) * 100
    
    axes[0, 1].semilogy(exceedance, q_sorted, color='blue', linewidth=2)
    axes[0, 1].set_xlabel('Exceedance Probability (%)', fontweight='bold')
    axes[0, 1].set_ylabel('Discharge (mm/day)', fontweight='bold')
    axes[0, 1].set_title('Flow Duration Curve', fontweight='bold')
    axes[0, 1].grid(True, alpha=0.3)
    
    # Add flow statistics
    q_stats = {
        'Q95': np.percentile(results['Q'], 5),
        'Q50': np.percentile(results['Q'], 50),
        'Q10': np.percentile(results['Q'], 90)
    }
    
    for stat_name, stat_value in q_stats.items():
        axes[0, 1].axhline(y=stat_value, linestyle='--', alpha=0.7, 
                          label=f'{stat_name}: {stat_value:.3f} mm/day')
    axes[0, 1].legend()
    
    # Runoff coefficient analysis
    runoff_coef = results['R'] / (P + 1e-6)
    runoff_coef[P < 0.1] = np.nan
    
    axes[1, 0].scatter(P, runoff_coef, alpha=0.6, c=results['W'], cmap='RdYlBu')
    axes[1, 0].set_xlabel('Precipitation (mm/day)', fontweight='bold')
    axes[1, 0].set_ylabel('Runoff Coefficient (-)', fontweight='bold')
    axes[1, 0].set_title('Runoff Coefficient vs Precipitation', fontweight='bold')
    axes[1, 0].set_ylim(0, 1)
    cbar3 = plt.colorbar(axes[1, 0].collections[0], ax=axes[1, 0])
    cbar3.set_label('Soil Moisture (mm)', fontweight='bold')
    axes[1, 0].grid(True, alpha=0.3)
    
    # Storage dynamics
    total_storage = results['W'] + results['S'] # Only use available storage components
    axes[1, 1].plot(dates[:len(total_storage)], total_storage, color='blue', linewidth=2, label='Total Storage')
    axes[1, 1].plot(dates[:len(results['W'])], results['W'], color='brown', linewidth=1.5, label='Soil Moisture', alpha=0.8)
    axes[1, 1].plot(dates[:len(results['S'])], results['S'], color='green', linewidth=1.5, label='Free Water', alpha=0.8)
    axes[1, 1].set_xlabel('Date', fontweight='bold')
    axes[1, 1].set_ylabel('Storage (mm)', fontweight='bold')
    axes[1, 1].set_title('Storage Components Dynamics', fontweight='bold')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)
    
    # Format x-axis for bottom plots
    axes[1, 1].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[1, 1].xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    plt.setp(axes[1, 1].xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'xinanjiang_performance.png'), dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Xinanjiang model visualization plots saved to '{save_dir}' directory")
    print("  Generated files:")
    print("  - xinanjiang_comprehensive.png")
    print("  - xinanjiang_water_balance.png")
    print("  - xinanjiang_performance.png")


def main():
    """
    Example usage of the Xinanjiang model with randomly generated data.
    """
    print("=" * 80)
    print("Xinanjiang (新安江) Hydrological Model - Example")
    print("=" * 80)
    
    # Set random seed for reproducibility
    np.random.seed(42)
    
    # Generate synthetic input data (365 days)
    n_days = 365
    
    # Precipitation: gamma distribution with shape=2, scale=5 (typical rainfall pattern)
    # Some days have no rain (60% dry days)
    P = np.random.gamma(2, 5, n_days)
    P[np.random.rand(n_days) > 0.4] = 0  # 60% dry days
    
    # Potential evapotranspiration: seasonal variation
    # Higher in summer, lower in winter
    t = np.arange(n_days)
    EP = 3.0 + 2.0 * np.sin(2 * np.pi * t / 365)  # Varies between 1-5 mm/day
    
def main():
    """
    Enhanced demonstration of Xinanjiang (新安江) Model with comprehensive analysis.
    """
    print("=" * 80)
    print("Xinanjiang (新安江) Hydrological Model - Enhanced Demonstration")
    print("Saturation Excess Runoff Generation Model for Humid Regions")
    print("=" * 80)
    
    # Set random seed for reproducibility
    np.random.seed(42)
    
    # Generate more realistic synthetic data (2.5 years)
    n_days = 912  # About 2.5 years
    
    # Create realistic precipitation pattern for humid climate
    t = np.arange(n_days)
    
    # Seasonal precipitation pattern (summer wet, winter dry)
    seasonal_factor = 1.8 + 1.2 * np.sin(2 * np.pi * t / 365 - np.pi/2)  # Peak in summer
    
    # Generate precipitation with mixed distributions
    P_base = np.random.gamma(2.0, 3.5, n_days) * seasonal_factor
    
    # Add monsoon-like wet periods
    wet_periods = np.random.choice(n_days, size=15, replace=False)
    for wp in wet_periods:
        duration = np.random.randint(5, 15)
        end_period = min(wp + duration, n_days)
        P_base[wp:end_period] *= np.random.uniform(2.0, 4.0)
    
    # Add dry periods (40% chance of no rain)
    dry_prob = 0.4 - 0.1 * np.sin(2 * np.pi * t / 365 - np.pi/2)  # Less dry in summer
    P = np.where(np.random.rand(n_days) < dry_prob, 0, P_base)
    
    # Add extreme events
    extreme_events = np.random.choice(n_days, size=12, replace=False)
    P[extreme_events] = P[extreme_events] + np.random.gamma(6, 12, 12)
    
    # Realistic potential evapotranspiration with seasonal variation
    EP_mean = 4.2
    EP_amplitude = 2.8
    EP = EP_mean + EP_amplitude * np.sin(2 * np.pi * t / 365) + np.random.normal(0, 0.25, n_days)
    EP = np.maximum(EP, 0.8)  # Minimum ET
    
    # Add temperature effect
    temp_effect = 1 + 0.4 * np.sin(2 * np.pi * t / 365)
    EP = EP * temp_effect
    
    print("\nEnhanced Input Data Summary:")
    print(f"  Simulation period: {n_days} days ({n_days/365:.1f} years)")
    print(f"  Total precipitation: {np.sum(P):.2f} mm")
    print(f"  Average daily precipitation: {np.mean(P):.2f} mm")
    print(f"  Maximum daily precipitation: {np.max(P):.2f} mm")
    print(f"  Rainfall days: {np.sum(P > 0.1)} ({np.sum(P > 0.1)/n_days*100:.1f}%)")
    print(f"  Total potential ET: {np.sum(EP):.2f} mm")
    print(f"  Average daily ET: {np.mean(EP):.2f} mm")
    print(f"  P/ET ratio: {np.sum(P)/np.sum(EP):.3f}")
    
    # Seasonal analysis
    seasons = ['Spring', 'Summer', 'Autumn', 'Winter']
    season_masks = [
        (t % 365 >= 60) & (t % 365 < 152),   # Spring (Mar-May)
        (t % 365 >= 152) & (t % 365 < 244),  # Summer (Jun-Aug)
        (t % 365 >= 244) & (t % 365 < 335),  # Autumn (Sep-Nov)
        (t % 365 >= 335) | (t % 365 < 60)    # Winter (Dec-Feb)
    ]
    
    print(f"\nSeasonal Distribution:")
    for season, mask in zip(seasons, season_masks):
        season_P = np.sum(P[mask])
        season_ET = np.sum(EP[mask])
        print(f"  {season}: P = {season_P:.1f} mm, ET = {season_ET:.1f} mm, P/ET = {season_P/season_ET:.2f}")
    
    # Initialize model with parameters suitable for humid regions
    model = XinanjiangModel(
        K=1.0,      # ET coefficient
        B=0.25,     # Tension water curve exponent
        IMP=0.02,   # Impervious area ratio
        WM=180.0,   # Average tension water capacity
        WUM=35.0,   # Upper layer capacity
        WLM=85.0,   # Lower layer capacity
        C=0.18,     # Deep layer ET coefficient
        SM=35.0,    # Free water capacity
        EX=1.2,     # Free water curve exponent
        KI=0.35,    # Interflow coefficient
        KG=0.25,    # Groundwater coefficient
        CI=0.75,    # Interflow recession
        CG=0.985    # Groundwater recession
    )
    
    print("\nEnhanced Model Parameters:")
    print(f"  Evapotranspiration:")
    print(f"    K (ET coefficient): {model.K}")
    print(f"    C (Deep layer ET coefficient): {model.C}")
    print(f"  Soil Moisture Structure:")
    print(f"    WM (Average tension water capacity): {model.WM} mm")
    print(f"    WUM (Upper layer capacity): {model.WUM} mm")
    print(f"    WLM (Lower layer capacity): {model.WLM} mm")
    print(f"    B (Tension water curve exponent): {model.B}")
    print(f"  Runoff Generation:")
    print(f"    SM (Free water capacity): {model.SM} mm")
    print(f"    EX (Free water curve exponent): {model.EX}")
    print(f"    IMP (Impervious area ratio): {model.IMP}")
    print(f"  Flow Routing:")
    print(f"    KI (Interflow coefficient): {model.KI}")
    print(f"    KG (Groundwater coefficient): {model.KG}")
    print(f"    CI (Interflow recession): {model.CI}")
    print(f"    CG (Groundwater recession): {model.CG}")
    
    # Run model
    print(f"\nRunning enhanced Xinanjiang model for {n_days} days...")
    results = model.run(P, EP)
    
    # Comprehensive results analysis
    print("\nComprehensive Simulation Results:")
    print(f"  Water Balance:")
    print(f"    Total Precipitation: {np.sum(P):.2f} mm")
    print(f"    Total Actual ET: {np.sum(results['E']):.2f} mm")
    print(f"    Total Runoff Generation: {np.sum(results['R']):.2f} mm")
    print(f"    Total Discharge: {np.sum(results['Q']):.2f} mm")
    print(f"    Annual Precipitation: {np.sum(P)*365/n_days:.2f} mm/year")
    print(f"    Annual Discharge: {np.sum(results['Q'])*365/n_days:.2f} mm/year")
    
    print(f"  Efficiency Indicators:")
    print(f"    Runoff Coefficient: {np.sum(results['R']) / np.sum(P):.3f}")
    print(f"    ET Efficiency: {np.sum(results['E']) / np.sum(EP):.3f}")
    print(f"    Water Use Efficiency: {np.sum(results['E']) / np.sum(P):.3f}")
    
    print(f"  Flow Characteristics:")
    print(f"    Peak discharge: {np.max(results['Q']):.2f} mm/day")
    print(f"    Mean discharge: {np.mean(results['Q']):.2f} mm/day")
    print(f"    Median discharge: {np.median(results['Q']):.2f} mm/day")
    print(f"    Q95 (low flows): {np.percentile(results['Q'], 5):.3f} mm/day")
    print(f"    Q10 (high flows): {np.percentile(results['Q'], 90):.2f} mm/day")
    print(f"    Flow variability (CV): {np.std(results['Q'])/np.mean(results['Q']):.3f}")
    
    print(f"  Soil Moisture Statistics:")
    print(f"    Average soil moisture: {np.mean(results['W']):.2f} mm ({np.mean(results['W'])/model.WM*100:.1f}% capacity)")
    print(f"    Maximum soil moisture: {np.max(results['W']):.2f} mm ({np.max(results['W'])/model.WM*100:.1f}% capacity)")
    print(f"    Minimum soil moisture: {np.min(results['W']):.2f} mm ({np.min(results['W'])/model.WM*100:.1f}% capacity)")
    print(f"    Soil moisture variability (CV): {np.std(results['W'])/np.mean(results['W']):.3f}")
    
    print(f"\nRunoff Components Analysis:")
    total_discharge = np.sum(results['Q'])
    surface_discharge = np.sum(results['QS'])
    interflow_discharge = np.sum(results['QI'])
    groundwater_discharge = np.sum(results['QG'])
    
    print(f"  Surface Runoff: {surface_discharge:.2f} mm ({surface_discharge/total_discharge*100:.1f}%)")
    print(f"  Interflow: {interflow_discharge:.2f} mm ({interflow_discharge/total_discharge*100:.1f}%)")
    print(f"  Groundwater: {groundwater_discharge:.2f} mm ({groundwater_discharge/total_discharge*100:.1f}%)")
    
    print(f"  Response Characteristics:")
    print(f"    Average surface runoff: {np.mean(results['QS']):.3f} mm/day")
    print(f"    Average interflow: {np.mean(results['QI']):.3f} mm/day")
    print(f"    Average groundwater flow: {np.mean(results['QG']):.3f} mm/day")
    
    # Water balance check
    total_input = np.sum(P)
    total_output = np.sum(results['E']) + np.sum(results['Q'])
    initial_storage = model.WUM/2 + model.WLM/2 + model.WM/3  # Initial estimate
    final_storage = results['W'][-1] + results['S'][-1]  # Only available storages
    storage_change = final_storage - initial_storage
    balance_error = total_input - total_output - storage_change
    
    print(f"\nWater Balance Check:")
    print(f"  Total input (P): {total_input:.2f} mm")
    print(f"  Total output (E + Q): {total_output:.2f} mm")
    print(f"  Storage change: {storage_change:.2f} mm")
    print(f"  Balance error: {balance_error:.2f} mm ({abs(balance_error)/total_input*100:.3f}%)")
    
    # Generate comprehensive visualizations
    print(f"\nGenerating comprehensive visualizations...")
    create_xinanjiang_plots(model, P, EP, results, save_dir="figures")
    
    # Enhanced sample output
    print("\nDetailed Results for First 20 Days:")
    print("Day |   Date   |   P    |   EP   |   E    |   Q    |   W    |   S    | QS   | QI   | QG   |")
    print("----|----------|--------|--------|--------|--------|--------|--------|------|------|------|")
    
    start_date = datetime(2020, 1, 1)
    for i in range(20):
        date_str = (start_date + timedelta(days=i)).strftime("%m-%d")
        print(f"{i+1:3d} | {date_str} | {P[i]:6.2f} | {EP[i]:6.2f} | {results['E'][i]:6.2f} | "
              f"{results['Q'][i]:6.2f} | {results['W'][i]:6.2f} | {results['S'][i]:6.2f} | "
              f"{results['QS'][i]:4.3f} | {results['QI'][i]:4.3f} | {results['QG'][i]:4.3f} |")
    
    # Model performance indicators
    print(f"\nModel Performance Indicators:")
    
    # Calculate some hydrological indices
    baseflow_index = groundwater_discharge / total_discharge
    quick_flow = surface_discharge + interflow_discharge
    quick_flow_index = quick_flow / total_discharge
    
    print(f"  Baseflow Index (BFI): {baseflow_index:.3f}")
    print(f"  Quickflow Index: {quick_flow_index:.3f}")
    print(f"  Regime characteristics: {'Baseflow dominated' if baseflow_index > 0.5 else 'Surface flow dominated'}")
    
    # Seasonal analysis of results
    print(f"\nSeasonal Performance Analysis:")
    for season, mask in zip(seasons, season_masks):
        if np.any(mask):
            season_Q = np.mean(results['Q'][mask])
            season_E = np.mean(results['E'][mask])
            season_W = np.mean(results['W'][mask])
            print(f"  {season}: Q = {season_Q:.2f} mm/d, E = {season_E:.2f} mm/d, W = {season_W:.1f} mm")
    
    print("\n" + "=" * 80)
    print("Enhanced Xinanjiang (新安江) simulation completed successfully!")
    print("This model demonstrates:")
    print("  - Saturation excess runoff generation mechanism")
    print("  - Three-layer soil moisture accounting")
    print("  - Three-component runoff separation")
    print("  - Suitable for humid and semi-humid regions")
    print("Check the 'figures' directory for comprehensive visualizations.")
    print("=" * 80)
    
    return results, model


if __name__ == "__main__":
    main()