#!/usr/bin/env python3
"""
Test para verificar la corrección de initial_head en TSNet
"""

import sys
import os
import tempfile


def test_tsnet_with_initial_head_fix():
    """
    Prueba TSNet con la corrección de initial_head
    """
    print("=== Test de corrección initial_head ===")

    try:
        import tsnet
        import wntr

        print(f"✓ TSNet version: {tsnet.__version__}")
        print(f"✓ WNTR version: {wntr.__version__}")
    except ImportError as e:
        print(f"✗ Error al importar: {e}")
        return False

    # Crear un archivo INP simple para testing
    inp_content = """[TITLE]
Test Network for initial_head fix

[JUNCTIONS]
;ID              	Elev        	Demand      	Pattern         
 J1              	0.00        	0.00      	;
 J2              	0.00        	0.00      	;

[RESERVOIRS]
;ID              	Head        	Pattern         
 R1              	10.00        	;

[TANKS]
;ID              	Elevation   	InitLevel   	MinLevel    	MaxLevel    	Diameter    	MinVol      	VolCurve
 T1              	0.00      	1.0         	0.0         	2.0         	1.0         	0           	;

[PIPES]
;ID              	Node1           	Node2           	Length      	Diameter    	Roughness   	MinorLoss   	Status
 P1              	R1              	J1              	100.0      	50.0       	0.001    	0           	Open
 P2              	J1              	J2              	100.0      	50.0       	0.001    	0           	Open
 P3              	J2              	T1              	100.0      	50.0       	0.001    	0           	Open

[PUMPS]
;ID              	Node1           	Node2           	Parameters

[VALVES]
;ID              	Node1           	Node2           	Diameter    	Type	Setting     	MinorLoss   

[PATTERNS]
;ID              	Multipliers
 1               	1.0    1.0    1.0    1.0    1.0    1.0    1.0    1.0
 1               	1.0    1.0    1.0    1.0    1.0    1.0    1.0    1.0
 1               	1.0    1.0    1.0    1.0    1.0    1.0    1.0    1.0

[CURVES]
;ID              	X-Value     	Y-Value

[QUALITY]
;Node            	InitQual
 J1              	0.0
 J2              	0.0
 R1              	0.0
 T1              	0.0

[OPTIONS]
 Units           	LPS
 Headloss        	D-W
 Specific Gravity	1.0
 Viscosity        	1.00E-003
 Trials          	40
 Accuracy        	0.001
 CHECKFREQ       	2
 MAXCHECK        	10
 DAMPLIMIT       	0
 Unbalanced      	Continue 10
 Pattern         	1
 Demand Multiplier	1.0
 Emitter Exponent	0.5
 Quality         	None mg/L
 Diffusivity     	1.0
 Tolerance       	0.01

[TIMES]
 Duration        	0:00:01
 Hydraulic Timestep	0:00:01
 Quality Timestep	0:05:00
 Pattern Timestep	1:00:00
 Pattern Start   	0:00:00
 Report Timestep 	1:00:00
 Report Start    	0:00:00
 Start ClockTime 	12:00:00 AM
 Statistic       	None

[COORDINATES]
;Node            	X-Coord         	Y-Coord
 J1              	100.00           	0.00
 J2              	200.00           	0.00
 R1              	0.00           	0.00
 T1              	300.00           	0.00

[END]
"""

    # Crear directorio y archivo temporal
    temp_dir = r"c:\Users\migue\AppData\Local\Temp\TSNet"
    os.makedirs(temp_dir, exist_ok=True)
    inp_path = os.path.join(temp_dir, "test_initial_head.inp")

    with open(inp_path, "w") as f:
        f.write(inp_content)

    print(f"✓ Archivo INP creado: {inp_path}")

    try:
        print("\n=== Aplicando corrección inicial_head (simulando PythonInterop) ===")

        # Cargar modelo con TSNet
        tm = tsnet.network.TransientModel(inp_path)
        print("✓ TSNet cargó el modelo")

        # Aplicar correcciones de división por cero
        if hasattr(tm, "simulation_period") and tm.simulation_period <= 0:
            tm.simulation_period = 1.0
            print("✓ simulation_period corregido = 1.0")

        if hasattr(tm, "time_step") and tm.time_step <= 0:
            tm.time_step = 0.1
            print("✓ time_step corregido = 0.1")

        # APLICAR CORRECCIÓN INITIAL_HEAD
        print("Aplicando corrección initial_head...")
        pipes_fixed = 0
        if hasattr(tm, "links"):
            for link in tm.links:
                # Verificar si es un Pipe y no tiene initial_head
                if "Pipe" in str(type(link)) and not hasattr(link, "initial_head"):
                    try:
                        # Asignar valor inicial por defecto
                        link.initial_head = 0.0
                        pipes_fixed += 1
                    except:
                        pass  # Si no se puede asignar, continuar

        print(f"✓ {pipes_fixed} pipes corregidos con initial_head")

        # Intentar simulación
        print("\n=== Intentando simulación TSNet ===")
        results = tsnet.simulation.MOCSimulator(
            tm, results_obj="results", friction="steady"
        )

        print("🎉 ¡SIMULACIÓN TSNET EXITOSA SIN FALLBACK!")
        print("✅ La corrección de initial_head funciona correctamente")

        return True

    except Exception as e:
        print(f"✗ Error: {e}")

        # Verificar si el error aún es sobre initial_head
        if "initial_head" in str(e):
            print("❌ La corrección de initial_head no fue suficiente")
            print("Se requiere investigación adicional")
        else:
            print("✅ El problema de initial_head está resuelto")
            print(f"Nuevo error diferente: {e}")

        return False


def main():
    """Función principal"""
    success = test_tsnet_with_initial_head_fix()

    if success:
        print("\n🎉 ¡CORRECCIÓN EXITOSA!")
        print("TSNet puede ejecutar simulaciones sin fallback a WNTR")
    else:
        print("\n⚠️ CORRECCIÓN PARCIAL")
        print("Verificar si el problema cambió o si se necesitan ajustes")

    return success


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)