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_dV_int, dV_int = find_intersections(t_common, dV_A_dt, dV_B_dt)

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

# ----------------------------
# Dimensionless scaling
# ----------------------------
V0_dimless = V_A_grid[-1]       # total volume
T_c = H0 / a_out * np.sqrt(1/(2*g))  # characteristic time

# Dimensionless variables
tA_nd = tA / T_c
tB_nd = tB / T_c
t_common_nd = t_common / T_c
hA_nd = hA_c / H0
hB_nd = hB_c / H0
dhA_nd = dhA_c * T_c / H0
dhB_nd = dhB_c * T_c / H0
V_A_nd = V_A / V0_dimless
V_B_nd = V_B / V0_dimless
dV_A_nd = dV_A_dt * T_c / V0_dimless
dV_B_nd = dV_B_dt * T_c / V0_dimless

# Dimensionless intersections
t_h_int_nd = t_h_int / T_c
t_dh_int_nd = t_dh_int / T_c
t_dV_int_nd = t_dV_int / T_c
dV_int_nd = dV_int * T_c / V0_dimless
h_dh_nd = h_dh_int / H0
dh_dh_nd = dh_h_int * T_c / H0

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

# Row 1: Volume-based
plt.subplot(2,3,1)
plt.plot(t_common_nd, V_A_nd, label='Tank A', color='blue')
plt.plot(t_common_nd, V_B_nd, label='Tank B', color='red')
plt.scatter([t_empty_A, t_empty_B]/T_c, [0,0], color=['blue','red'])
plt.text(t_empty_A/T_c,0,f"{t_empty_A/T_c:.2f}", color='blue', va='top')
plt.text(t_empty_B/T_c,0,f"{t_empty_B/T_c:.2f}", color='red', va='top')
plt.xlabel(r"$\overline{t}=t/T_c$"); plt.ylabel(r"$\overline{V}=V/V_0$")
plt.grid(True); plt.legend()

plt.subplot(2,3,2)
plt.plot(hA_nd[::-1], V_A_nd[::-1], label='Tank A', color='blue')
plt.plot(hB_nd[::-1], V_B_nd[::-1], label='Tank B', color='red')
plt.xlabel(r"$\overline{h}=h/H_0$"); plt.ylabel(r"$\overline{V}=V/V_0$")
plt.gca().invert_xaxis()
plt.grid(True); plt.legend()

plt.subplot(2,3,3)
plt.plot(t_common_nd, dV_A_nd, label='Tank A', color='blue')
plt.plot(t_common_nd, dV_B_nd, label='Tank B', color='red')
for i in range(len(t_dV_int_nd)):
    plt.scatter(t_dV_int_nd[i], dV_int_nd[i], color='purple')
    plt.text(t_dV_int_nd[i], dV_int_nd[i], f"({t_dV_int_nd[i]:.2f},{dV_int_nd[i]:.4f})", color='purple')
plt.xlabel(r"$\overline{t}=t/T_c$"); plt.ylabel(r"$d\overline{V}/d\overline{t}$")
plt.grid(True); plt.legend()

# Row 2: Height-based
plt.subplot(2,3,4)
plt.plot(t_common_nd, hA_nd, label='Tank A', color='blue')
plt.plot(t_common_nd, hB_nd, label='Tank B', color='red')
plt.scatter([t_empty_A, t_empty_B]/T_c, [0,0], color=['blue','red'])
plt.text(t_empty_A/T_c,0,f"{t_empty_A/T_c:.2f}", color='blue', va='top')
plt.text(t_empty_B/T_c,0,f"{t_empty_B/T_c:.2f}", color='red', va='top')
plt.xlabel(r"$\overline{t}=t/T_c$"); plt.ylabel(r"$\overline{h}=h/H_0$")
plt.grid(True); plt.legend()
for i in range(len(t_h_int_nd)):
    plt.scatter(t_h_int_nd[i], h_int[i]/H0, color='purple')
    plt.text(t_h_int_nd[i], h_int[i]/H0, f"({t_h_int_nd[i]:.2f},{h_int[i]/H0:.2f})", color='purple')

plt.subplot(2,3,5)
plt.plot(t_common_nd, dhA_nd, label='Tank A', color='blue')
plt.plot(t_common_nd, dhB_nd, label='Tank B', color='red')
plt.xlabel(r"$\overline{t}=t/T_c$"); plt.ylabel(r"$d\overline{h}/d\overline{t}$")
plt.grid(True); plt.legend()
for i in range(len(t_dh_int_nd)):
    plt.scatter(t_dh_int_nd[i], dh_int[i]*T_c/H0, color='purple')
    plt.text(t_dh_int_nd[i], dh_int[i]*T_c/H0, f"({t_dh_int_nd[i]:.2f},{dh_int[i]*T_c/H0:.4f})", color='purple')

plt.subplot(2,3,6)
plt.plot(h_grid/H0, dhA_h * T_c/H0, label='Tank A', color='blue')
plt.plot(h_grid/H0, dhB_h * T_c/H0, label='Tank B', color='red')
for i in range(len(h_dh_nd)):
    plt.scatter(h_dh_nd[i], dh_dh_nd[i], color='purple')
    plt.text(h_dh_nd[i], dh_dh_nd[i], f"({h_dh_nd[i]:.2f},{dh_dh_nd[i]:.4f})", color='purple')
plt.xlabel(r"$\overline{h}=h/H_0$"); plt.ylabel(r"$d\overline{h}/d\overline{t}$")
plt.gca().invert_xaxis()
plt.grid(True); plt.legend()

plt.tight_layout()
plt.show()