#!/usr/bin/env python3
"""
Before/After Comparison Test
Tests the difference between nstp=[1] and nstp=[10] to show non-linear improvement
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from gspy_interface import run_modflow_step

def test_linear_vs_nonlinear():
    """Compare linear (old) vs non-linear (new) behavior"""
    
    print("=" * 80)
    print("BEFORE/AFTER TIME DISCRETIZATION COMPARISON")
    print("Demonstrates improvement from nstp=[1] to nstp=[10]")
    print("=" * 80)
    
    # Test Period 3 pumping behavior
    print("\nTesting Period 3 pumping at -500 m³/day:")
    print("HK=10.0, Sy=0.3 (Tutorial 2 standard parameters)")
    print("\nExpected with nstp=[10]: Non-linear drawdown curve")
    print("Previous with nstp=[1]: Linear drawdown")
    print()
    
    # Reset state for clean test
    from state_manager import StateManager
    state_mgr = StateManager()
    state_mgr.reset_state()
    print("✓ State reset for clean test")
    
    # Run baseline periods
    print("\n--- Establishing Baseline ---")
    baseline = run_modflow_step(1, 0.0, 10.0, 0.3)[0]  # Period 1
    print(f"Period 1 baseline: {baseline:.6f} m")
    
    # Period 2 for gradient
    for day in range(2, 5):
        head = run_modflow_step(2, 0.0, 10.0, 0.3)[0]
        print(f"Period 2 Day {day}: {head:.6f} m")
    
    print("\n--- Period 3 Pumping Analysis ---")
    print("Day   Head (m)     Δ Head      Rate Change  Behavior")
    print("-" * 55)
    
    prev_head = None
    prev_rate = None
    heads = []
    
    for day in range(5, 15):  # 10 days of pumping
        head = run_modflow_step(3, -500.0, 10.0, 0.3)[0]
        heads.append(head)
        
        if prev_head is not None:
            delta_head = head - prev_head
            rate = abs(delta_head)
            
            behavior = ""
            if prev_rate is not None:
                rate_change = rate - prev_rate
                if abs(rate_change) < 0.000001:
                    behavior = "LINEAR"
                elif rate_change > 0:
                    behavior = "ACCELERATING"
                else:
                    behavior = "DECELERATING"
            
            print(f"{day:3d}   {head:.6f}   {delta_head:8.6f}   {rate_change if prev_rate else '   --    ':9.6f}  {behavior}")
            prev_rate = rate
        else:
            print(f"{day:3d}   {head:.6f}      --         --       --")
        
        prev_head = head
    
    print("\n--- Analysis Summary ---")
    
    # Calculate behavior metrics
    early_heads = heads[1:4]  # Days 6-8
    late_heads = heads[7:10]  # Days 12-14
    
    early_avg = sum(early_heads) / len(early_heads)
    late_avg = sum(late_heads) / len(late_heads)
    
    total_drawdown = heads[0] - heads[-1]
    avg_rate = total_drawdown / len(heads)
    
    print(f"Total drawdown over {len(heads)} days: {total_drawdown:.6f} m")
    print(f"Average drawdown rate: {avg_rate:.6f} m/day")
    print(f"Early period average head: {early_avg:.6f} m")
    print(f"Late period average head: {late_avg:.6f} m")
    
    # Check for non-linearity
    rates = []
    for i in range(1, len(heads)):
        rate = abs(heads[i] - heads[i-1])
        rates.append(rate)
    
    rate_variation = max(rates) - min(rates) if rates else 0
    print(f"Rate variation: {rate_variation:.8f} m/day")
    
    if rate_variation < 0.000010:
        print("⚠ NEARLY LINEAR behavior detected")
        print("  This suggests daily time steps are still too coarse")
        print("  Consider shorter time steps or different model parameters")
    else:
        print("✓ NON-LINEAR behavior detected")
        print("  Transient dynamics are being captured properly")
    
    print("\n" + "=" * 80)
    print("COMPARISON COMPLETE")
    print("Current implementation uses nstp=[10] for improved transient capture")
    print("Previous nstp=[1] approach produced linear drawdown")
    print("=" * 80)

if __name__ == "__main__":
    test_linear_vs_nonlinear()