import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp, cumulative_trapezoid
from scipy.interpolate import interp1d

# ----------------------------
# Physical parameters
# ----------------------------
g = 9.81
a_out = 0.01
H0 = 1.0

# ----------------------------
# Tank volume scaling
# ----------------------------
V_target = 1.0
r = 4
A_small = 3*V_target / (H0 * (1 + r + np.sqrt(r)))
A_large = r * A_small

# ----------------------------
# Tank geometries
# ----------------------------
def A_A(h):
    return A_small + (A_large - A_small) * h / H0

def A_B(h):
    return A_large - (A_large - A_small) * h / H0

# ----------------------------
# dh/dt ODEs
# ----------------------------
def dhdt_A(t, h):
    return -a_out * np.sqrt(2*g*h[0]) / A_A(h[0]) if h[0] > 0 else 0

def dhdt_B(t, h):
    return -a_out * np.sqrt(2*g*h[0]) / A_B(h[0]) if h[0] > 0 else 0

def empty_event(t, h):
    return h[0]
empty_event.terminal = True
empty_event.direction = -1

# ----------------------------
# Solve ODEs
# ----------------------------
solA = solve_ivp(dhdt_A, [0, 500], [H0], events=empty_event, max_step=0.01)
solB = solve_ivp(dhdt_B, [0, 500], [H0], events=empty_event, max_step=0.01)

tA, hA = solA.t, solA.y[0]
tB, hB = solB.t, solB.y[0]

t_final = max(tA[-1], tB[-1])
t_common = np.linspace(0, t_final, 8000)

hA_c = interp1d(tA, hA, kind='cubic', fill_value=0, bounds_error=False)(t_common)
hB_c = interp1d(tB, hB, kind='cubic', fill_value=0, bounds_error=False)(t_common)

dhA_c = -a_out * np.sqrt(2*g*np.maximum(hA_c,1e-6)) / A_A(np.maximum(hA_c,1e-6))
dhB_c = -a_out * np.sqrt(2*g*np.maximum(hB_c,1e-6)) / A_B(np.maximum(hB_c,1e-6))

# ----------------------------
# Volume vs Height
# ----------------------------
h_grid_fine = np.linspace(0, H0, 2000)
V_A_grid = cumulative_trapezoid(A_A(h_grid_fine), h_grid_fine, initial=0)
V_B_grid = cumulative_trapezoid(A_B(h_grid_fine), h_grid_fine, initial=0)

V_A = np.interp(hA_c, h_grid_fine, V_A_grid)
V_B = np.interp(hB_c, h_grid_fine, V_B_grid)

dV_A_dt = dhA_c * A_A(hA_c)
dV_B_dt = dhB_c * A_B(hB_c)

h_grid = np.linspace(0.0001, H0, 8000)
dhA_h = -a_out * np.sqrt(2*g*h_grid)/A_A(h_grid)
dhB_h = -a_out * np.sqrt(2*g*h_grid)/A_B(h_grid)

# ----------------------------
# Find intersections
# ----------------------------
def find_intersections(x, y1, y2):
    diff = y1 - y2
    sign_change = np.where(np.diff(np.sign(diff)) != 0)[0]
    xs, ys = [], []
    for i in sign_change:
        x0, x1 = x[i], x[i+1]
        y0, y1_ = diff[i], diff[i+1]
        x_root = x0 - y0*(x1-x0)/(y1_-y0)
        y_root = np.interp(x_root, x, y1)
        xs.append(x_root)
        ys.append(y_root)
    return np.array(xs), np.array(ys)

# Intersections
t_h_int, h_int = find_intersections(t_common, hA_c, hB_c)
if len(t_h_int) > 0:
    t_h_int = t_h_int[1:]
    h_int = h_int[1:]

t_dh_int, dh_int = find_intersections(t_common, dhA_c, dhB_c)
h_dh_int, dh_h_int = find_intersections(h_grid, dhA_h, dhB_h)

t_empty_A, t_empty_B = tA[-1], tB[-1]
h_empty = 0.0

# ----------------------------
# Plotting 6 subplots
# ----------------------------
plt.figure(figsize=(18,10))

# Row 1: Volume-based
plt.subplot(2,3,1)
plt.plot(t_common, V_A, label='Tank A', color='blue')
plt.plot(t_common, V_B, label='Tank B', color='red')
plt.scatter([t_empty_A, t_empty_B], [0,0], color=['blue','red'])
plt.text(t_empty_A,0,f"{t_empty_A:.2f}", color='blue', va='top')
plt.text(t_empty_B,0,f"{t_empty_B:.2f}", color='red', va='top')
plt.xlabel("Time [s]"); plt.ylabel("Volume [m³]"); plt.title("Volume vs Time")
plt.grid(True); plt.legend()

plt.subplot(2,3,2)
plt.plot(hA_c[::-1], V_A[::-1], label='Tank A', color='blue')
plt.plot(hB_c[::-1], V_B[::-1], label='Tank B', color='red')
plt.xlabel("Height [m]"); plt.ylabel("Volume [m³]"); plt.title("Volume vs Height")
plt.gca().invert_xaxis()  # descending x-axis
plt.grid(True); plt.legend()

plt.subplot(2,3,3)
plt.plot(t_common, dV_A_dt, label='Tank A', color='blue')
plt.plot(t_common, dV_B_dt, label='Tank B', color='red')

# Add intersections for dV/dt vs Time
t_dV_int, dV_int = find_intersections(t_common, dV_A_dt, dV_B_dt)
for i in range(len(t_dV_int)):
    plt.scatter(t_dV_int[i], dV_int[i], color='purple')
    plt.text(t_dV_int[i], dV_int[i], f"({t_dV_int[i]:.2f},{dV_int[i]:.4f})", color='purple')

plt.xlabel("Time [s]")
plt.ylabel("dV/dt [m³/s]")
plt.title("dV/dt vs Time")
plt.grid(True)
plt.legend()

# Row 2: Height-based
plt.subplot(2,3,4)
plt.plot(t_common, hA_c, label='Tank A', color='blue')
plt.plot(t_common, hB_c, label='Tank B', color='red')
plt.scatter([t_empty_A, t_empty_B], [h_empty,h_empty], color=['blue','red'])
plt.text(t_empty_A,0,f"{t_empty_A:.2f}", color='blue', va='top')
plt.text(t_empty_B,0,f"{t_empty_B:.2f}", color='red', va='top')
plt.xlabel("Time [s]"); plt.ylabel("Height [m]"); plt.title("Height vs Time")
plt.grid(True); plt.legend()
for i in range(len(t_h_int)):
    plt.scatter(t_h_int[i], h_int[i], color='purple')
    plt.text(t_h_int[i], h_int[i], f"({t_h_int[i]:.2f},{h_int[i]:.2f})", color='purple')

plt.subplot(2,3,5)
plt.plot(t_common, dhA_c, label='Tank A', color='blue')
plt.plot(t_common, dhB_c, label='Tank B', color='red')
plt.xlabel("Time [s]"); plt.ylabel("dh/dt [m/s]"); plt.title("dh/dt vs Time")
plt.grid(True); plt.legend()
for i in range(len(t_dh_int)-1):
    plt.scatter(t_dh_int[i], dh_int[i], color='purple')
    plt.text(t_dh_int[i], dh_int[i], f"({t_dh_int[i]:.2f},{dh_int[i]:.4f})", color='purple')

plt.subplot(2,3,6)
plt.plot(h_grid, dhA_h, label='Tank A', color='blue')
plt.plot(h_grid, dhB_h, label='Tank B', color='red')
plt.xlabel("Height [m]"); plt.ylabel("dh/dt [m/s]"); plt.title("dh/dt vs Height")
plt.gca().invert_xaxis()  # descending x-axis
plt.grid(True); plt.legend()
for i in range(len(h_dh_int)):
    plt.scatter(h_dh_int[i], dh_h_int[i], color='purple')
    plt.text(h_dh_int[i], dh_h_int[i], f"({h_dh_int[i]:.2f},{dh_h_int[i]:.4f})", color='purple')

plt.tight_layout()
plt.show()