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        # gravity m/s²
a_out = 0.01    # outlet area m²
H0 = 1.0        # tank height m

# ----------------------------
# Tank volume scaling
# ----------------------------
V_target = 1.0   # m³
r = 4            # ratio A_large / A_small

# Solve for small and large area
A_small = 3*V_target / (H0 * (1 + r + np.sqrt(r)))
A_large = r * A_small

# ----------------------------
# Tank geometries
# ----------------------------
def A_A(h):
    # Tank A: wide top -> narrow bottom
    return A_small + (A_large - A_small) * h / H0

def A_B(h):
    # Tank B: narrow top -> wide bottom
    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]

# ----------------------------
# Common time grid until both tanks empty
# ----------------------------
t_final = max(tA[-1], tB[-1])
t_common = np.linspace(0, t_final, 8000)

# Interpolate heights
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)

# dh/dt arrays
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))

# ----------------------------
# 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)

# Height vs Time intersections (skip first obvious at t=0)
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:]

# dh/dt vs Time intersections
t_dh_int, dh_int = find_intersections(t_common, dhA_c, dhB_c)

# dh/dt vs Height intersections
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)
h_dh_int, dh_h_int = find_intersections(h_grid, dhA_h, dhB_h)

# ----------------------------
# Volume vs Time
# ----------------------------
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)

# Emptying times for points
t_empty_A = tA[-1]
t_empty_B = tB[-1]
h_empty = 0.0

# ----------------------------
# Plotting
# ----------------------------
plt.figure(figsize=(14,10))

# Volume vs Time
plt.subplot(2,2,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'A empty t={t_empty_A:.2f}s',color='blue', va='top')
plt.text(t_empty_B,0,f'B empty t={t_empty_B:.2f}s',color='red', va='top')
plt.grid(True)
plt.xlabel("Time [s]")
plt.ylabel("Volume [m³]")
plt.title("Volume vs Time")
plt.legend()

# Height vs Time
plt.subplot(2,2,2)
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,h_empty,f'A empty t={t_empty_A:.2f}s',color='red', va='top')
plt.text(t_empty_B,h_empty,f'B empty t={t_empty_B:.2f}s',color='red', va='top')
for i in range(len(t_h_int)):
    plt.scatter(t_h_int[i], h_int[i], color='green')
    plt.text(t_h_int[i], h_int[i], f"({t_h_int[i]:.3f},{h_int[i]:.3f})", color='green')
plt.grid(True)
plt.xlabel("Time [s]")
plt.ylabel("Height [m]")
plt.title("Height vs Time")
plt.legend()

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

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

plt.tight_layout()
#plt.savefig("tank emptying comparison.png", dpi)
plt.show()

# ----------------------------
# Print intersections numerically
# ----------------------------
print("Height vs Time intersections (t,h):")
for i in range(len(t_h_int)):
    print(f"  ({i+1}) t={t_h_int[i]:.5f}, h={h_int[i]:.5f}")

print("\ndh/dt vs Time intersections (t,dh/dt):")
for i in range(len(t_dh_int)):
    print(f"  ({i+1}) t={t_dh_int[i]:.5f}, dh/dt={dh_int[i]:.5f}")

print("\ndh/dt vs h intersections (h,dh/dt):")
for i in range(len(h_dh_int)):
    print(f"  ({i+1}) h={h_dh_int[i]:.5f}, dh/dt={dh_h_int[i]:.5f}")