import torch
import matplotlib.pyplot as plt

# --- Parameters ---
nx = 100          # mesh points
dx = 1.0 / (nx-1)
dt = 0.0001       # time step
alpha = 0.01      # thermal diffusivity
steps = 500       # number of time steps

# --- Create mesh and initial condition (base state) ---
x = torch.linspace(0, 1, nx)
u_base = torch.sin(torch.pi * x)   # smooth initial profile

# --- Fixed boundary conditions (Dirichlet) ---
def apply_bc(u):
    u[0] = 0.0
    u[-1] = 0.0
    return u

# --- Time evolution using finite differences ---
def heat_solve(u0, steps, dt, alpha, dx):
    u = u0.clone()
    for _ in range(steps):
        u = apply_bc(u)
        u[1:-1] = u[1:-1] + alpha * dt / dx**2 * (u[2:] - 2*u[1:-1] + u[:-2])
    return u

# 1. Compute the "base finished state" (idling engine)
u_final_base = heat_solve(u_base, steps, dt, alpha, dx)

# 2. New "prompt": change initial condition to something else (e.g., a different sine)
u_new_initial = torch.sin(2 * torch.pi * x)

# 3. Cold start: solve from zero (or random) guess
u_cold = heat_solve(torch.zeros_like(x), steps, dt, alpha, dx)

# 4. Warm start: use the previous final state as initial guess
u_warm = heat_solve(u_final_base, steps, dt, alpha, dx)   # reusing u_final_base

# --- Compare errors (solutions should converge to same steady state) ---
print("Max difference between cold and warm solutions:",
      torch.max(torch.abs(u_cold - u_warm)).item())

# Visualize
plt.plot(x.numpy(), u_final_base.numpy(), label='Base state (old BC)')
plt.plot(x.numpy(), u_cold.numpy(), '--', label='Cold start (new BC)')
plt.plot(x.numpy(), u_warm.numpy(), ':', label='Warm start (reused state)')
plt.legend()
plt.show()