# -*- coding: utf-8 -*-
"""
Created on Tue Dec  3 10:59:59 2024

@author: Bijan
"""

import random
import numpy as np
import matplotlib.pyplot as plt

Lx, Ly = 1.0, 1.0  # Domain size (meters)
Nx, Ny = 20, 20    # Number of grid points
dx, dy = Lx / (Nx - 1), Ly / (Ny - 1)  # Grid spacing
alpha = 0.01       # Thermal diffusivity (m^2/s)
h2=min(dx,dy)**2
dt = 0.25*h2/alpha
T = 1            # Total simulation time (seconds)

# Define the grid
x = np.linspace(0, Lx, Nx)
y = np.linspace(0, Ly, Ny)
X, Y = np.meshgrid(x, y)

def plot_sol(u,label,title,xlabel,ylabel):
    plt.figure(figsize=(8, 6))
    plt.contourf(X, Y, u, 20, cmap='hot')
    plt.colorbar(label=label)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.show()

def numerical_integration(u,v):
    integral=0
    for i in range(Nx):
        for j in range(Ny):
            integral+=dx*dy*u[i,j]*v[i,j]    
    return integral 
    
def state_equation(source_strength,source_localization):
    # Parameters

    nb_ctrl=len(source_strength)
        
    # Initial condition: u = 0 everywhere
    u = np.zeros((Nx, Ny))
    
    source=np.zeros((Nx, Ny))
    for ictrl in range(nb_ctrl):
        x0, y0 = source_localization[ictrl,:]
        sigma = 0.01 #*(ictrl+1)            # Source width
        source += source_strength[ictrl] * np.exp(-((X - x0)**2 + (Y - y0)**2) / (2 * sigma**2))
    
    # Time-stepping loop
    nt = int(T / dt)  # Number of time steps
    u_new = np.copy(u)
    
    for n in range(nt):
        for i in range(1, Nx-1):
            for j in range(1, Ny-1):
                laplacian = (u[i+1, j] - 2*u[i, j] + u[i-1, j]) / dx**2 + \
                            (u[i, j+1] - 2*u[i, j] + u[i, j-1]) / dy**2
                u_new[i, j] = u[i, j] + dt * (alpha * laplacian + source[i,j])
        
        # Update the temperature array
        u = np.copy(u_new)

    return u
#%%

random.seed(0)

nb_ctrl=5
source_strength_target=np.arange(nb_ctrl)+1    
source_localization=np.zeros((nb_ctrl,2))
for ictrl in range(nb_ctrl):
    source_localization[ictrl,0]=random.random()
    source_localization[ictrl,1]=random.random()

u_target=state_equation(source_strength_target,source_localization)
plot_sol(u_target,"Temperature","Target","x","y")

source_strength=np.zeros(nb_ctrl)
u_elem_0=state_equation(source_strength,source_localization)

u_elem=[]
for ictrl in range(nb_ctrl):
    source_strength=np.zeros(nb_ctrl)
    source_strength[ictrl]=1
    u=state_equation(source_strength,source_localization)
    u_elem.append(u)
    #plot_sol(u,"label","title","xlabel","ylabel")

u_elem=np.array(u_elem)

A=np.zeros((nb_ctrl,nb_ctrl))
b=np.zeros(nb_ctrl)
for ictrl in range(nb_ctrl):
    b[ictrl]=numerical_integration(u_elem[ictrl,:,:], (u_target-u_elem_0))
    for jctrl in range(ictrl,nb_ctrl):
        A[ictrl,jctrl]=numerical_integration(u_elem[ictrl,:,:], u_elem[jctrl,:,:])
    for jctrl in range(0,ictrl):
        A[ictrl,jctrl]=A[jctrl,ictrl]

source_strength_opt = np.linalg.solve(A, b)

cost=numerical_integration(u_target-u, u_target-u)
print("Cost (error in state variable):",cost)
print("Error in control :",np.sum(np.abs(source_strength_opt-source_strength_target)))

u_opt=state_equation(source_strength_opt,source_localization)

plot_sol(u_target-u_opt,"Error","Target-Optimized","x","y")
print("sol opt inversion:",source_strength_opt)

#%%
from scipy.optimize import minimize
def cost(x):
    u=state_equation(x,source_localization)
    functional=numerical_integration(u_target-u, u_target-u)
    convergence.append(functional)
    return functional

convergence=[]
x0 = np.zeros(nb_ctrl)
x0=np.ones(nb_ctrl)+10
res = minimize(cost, x0, 
               method='Nelder-Mead', 
               options={'xatol': 1e-13, 
                        'disp': True})
print(res.fun)
print("sol opt direct:",res.x)

plt.plot(np.log10(convergence),label="Convergence Nelder-Mead")
plt.xlabel("Functional evaluations")
plt.legend()
plt.show()

u_opt=state_equation(res.x,source_localization)
plot_sol(u_target-u_opt,"Error","Target-Optimized","x","y")