#!/usr/bin/env python3
"""
Head Distribution Analysis
=========================
Analyze the complete head distribution to understand why
interior cells are not responding to boundary gradients.
"""

import os
import numpy as np
import flopy
from gspy_interface import run_modflow_step

def analyze_head_distribution():
    """Analyze the complete head distribution"""
    
    print("="*80)
    print("HEAD DISTRIBUTION ANALYSIS")
    print("Examining the complete head field")
    print("="*80)
    
    # Clean slate
    state_file = "modflow_state.pkl"
    if os.path.exists(state_file):
        os.remove(state_file)
    
    # Run Period 1 and 2
    print("Running Period 1 (uniform)...")
    run_modflow_step(1, 0.0, 10.0, 0.3)
    
    print("Running Period 2 (gradient)...")
    run_modflow_step(2, 0.0, 10.0, 0.3)
    
    # Read the head file directly
    head_file = "gspy_tutorial2.hds"
    if not os.path.exists(head_file):
        print(f"❌ Head file {head_file} not found!")
        return
    
    print(f"\n📊 Reading head file: {head_file}")
    
    try:
        # Read heads with FloPy
        hds = flopy.utils.HeadFile(head_file)
        times = hds.get_times()
        
        print(f"Available times: {times}")
        
        # Get the latest heads
        heads = hds.get_data(totim=times[-1])
        
        print(f"\n🗺️  COMPLETE HEAD DISTRIBUTION (Layer 0):")
        print("    Col:  0      1      2      3      4      5      6      7      8      9")
        print("         " + "─" * 70)
        
        for row in range(10):
            row_str = f"Row {row}: "
            for col in range(10):
                head_val = heads[0, row, col]
                row_str += f"{head_val:6.2f} "
            print(row_str)
        
        print("\n📈 GRADIENT ANALYSIS:")
        center_row = 4
        print(f"Center row ({center_row}): Left→Center→Right")
        
        for col in range(10):
            head = heads[0, center_row, col]
            if col == 0:
                print(f"  Col {col} (LEFT):   {head:7.3f} m")
            elif col == 4:
                print(f"  Col {col} (CENTER): {head:7.3f} m ← TARGET")
            elif col == 9:
                print(f"  Col {col} (RIGHT):  {head:7.3f} m")
            else:
                print(f"  Col {col}:         {head:7.3f} m")
        
        # Check for gradient
        left_boundary = heads[0, center_row, 0]
        right_boundary = heads[0, center_row, 9] 
        center_cell = heads[0, center_row, 4]
        
        expected_center = left_boundary + (right_boundary - left_boundary) * 0.4  # 40% across
        
        print(f"\n🎯 GRADIENT VERIFICATION:")
        print(f"Left boundary:  {left_boundary:.6f} m")
        print(f"Right boundary: {right_boundary:.6f} m")
        print(f"Center cell:    {center_cell:.6f} m")
        print(f"Expected center: {expected_center:.6f} m")
        print(f"Difference:     {abs(center_cell - expected_center):.6f} m")
        
        if abs(center_cell - expected_center) > 1.0:
            print("❌ MAJOR ISSUE: Center cell does not reflect gradient!")
        elif abs(center_cell - expected_center) > 0.1:
            print("⚠️  Minor deviation from expected gradient")
        else:
            print("✅ Center cell reflects proper gradient")
            
    except Exception as e:
        print(f"❌ Error reading head file: {e}")

if __name__ == "__main__":
    analyze_head_distribution()