# case1.py

import numpy as np
import matplotlib.pyplot as plt

from Bar import Bar, set_k_ratio

def get_analytical_constants(bar1:'Bar', bar2:'Bar', x_bar=2/np.pi, l=1):
    '''Solve for C1, D1, C2, and D2
    Returns np.array([C1, D1, C2, D2])'''
    epsilon = (bar1.k*bar1.area*bar1.alpha) / (bar2.k*bar2.area*bar2.alpha)
    A = np.array([
        #C1                                     D1              C2                              D2
        [0,                                     1,              0,                              0],
        [0,                                     0,              np.sinh(bar2.alpha*l),          np.cosh(bar2.alpha*l)],
        [np.sinh(bar1.alpha*x_bar),              0,             -1*np.sinh(bar2.alpha*x_bar),   -1*np.cosh(bar2.alpha*x_bar)],
        [epsilon*np.cosh(bar1.alpha*x_bar),      0,             -1*np.cosh(bar2.alpha*x_bar),   -1*np.sinh(bar2.alpha*x_bar)]
    ])

    B = np.array([
        0,
        100,
        0,
        0
    ])

    return np.linalg.solve(A, B)

def get_analytical_datapoints(x_points, bar1:'Bar', bar2:'Bar', x_bar=2/np.pi, l=1):
    '''Generates the temperature distribution given x_points.
    Returns np.array of tempature values'''
    constants = get_analytical_constants(bar1=bar1,bar2=bar2,x_bar=x_bar,l=1)
    C1, D1, C2, D2 = constants[0], constants[1], constants[2], constants[3]

    # Two temp functions for the left and right side of the interface
    left_temp = lambda x: C1*np.sinh(bar1.alpha*x) + D1*np.cosh(bar1.alpha*x)
    right_temp = lambda x: C2*np.sinh(bar2.alpha*x) + D2*np.cosh(bar2.alpha*x)

    temp_values = np.array([])
    for x in x_points:
        if x <= x_bar:
            temp = left_temp(x)
        else:
            temp = right_temp(x)
        temp_values = np.append(temp_values, temp)

    return temp_values

def get_analytical_heat_convection(x, bar1:'Bar', bar2:'Bar', x_bar=2/np.pi, l=1):
    '''Gets the analytical heat convection at x from heat flux conservation'''
    constants = get_analytical_constants(bar1=bar1,bar2=bar2,x_bar=x_bar,l=1)
    C1, D1, C2, D2 = constants[0], constants[1], constants[2], constants[3]

    t_prime = C2*bar2.alpha*np.cosh(bar2.alpha*x) + D2*bar2.alpha*np.sinh(bar2.alpha*x)

    q_dot = -bar2.k*bar2.area*t_prime
    return q_dot

def fdm_array(bar1:'Bar', bar2:'Bar', num_sections: int=6, x_bar=2/np.pi, l=1):
    '''num_sections = total sections to cut the bar. Note that this is NOT evenly distributed over the entire
    bar. Half of the sections are left of x_bar, half are to the right.
    Returns the A matrix of the FDM equation Ax = B'''
    dx1 = x_bar / (num_sections / 2)
    dx2 =  (l - x_bar) / (num_sections / 2)
    k1 = 2 + bar1.alpha**2 * dx1**2
    k2 = 2 + bar2.alpha**2 * dx2**2

    beta1 = (bar1.k * bar1.area) / (dx1)
    beta2 = (bar2.k * bar2.area) / (dx2)

    #rho = beta1 - beta2 - (bar1.k * bar1.area * dx1 * bar1.alpha**2) / 2 + (bar2.k * bar2.area * dx2 * bar2.alpha**2) / 2
    rho = beta1 + beta2 + (bar1.k * bar1.area * dx1 * bar1.alpha**2) / 2 + (bar2.k * bar2.area * dx2 * bar2.alpha**2) / 2
    # Making the matrix
    # There are num_sections - 1 rows and columns
    # row 0, row n - 2, and row num_sections/2 - 1 are unique
    # the rest are -1, k, -1 where k changes from k1 to k2 at num_sections/2

    # Initialize matrix A with zeros
    A = np.zeros((num_sections - 1, num_sections - 1))
    
    # Fill the matrix A
    for row in range(num_sections - 1):
        if row == 0:
            # First row for boundary at x = 0 (for example Dirichlet or Neumann condition)
            A[row, row] = k1
            A[row, row+1] = -1
        elif row == num_sections // 2 - 1:
            # Special row at the interface between bar1 and bar2
            A[row, row-1] = -beta1
            A[row, row] = rho  # This is the rho value you computed
            A[row, row+1] = -beta2
        elif row == num_sections - 2:
            # Last row for boundary at x = l (again assuming Dirichlet or Neumann condition)
            A[row, row-1] = -1
            A[row, row] = k2
        else:
            # Interior rows for bar1 and bar2 (uniform grid)
            if row < num_sections // 2 - 1:
                # In bar1 region
                A[row, row-1] = -1
                A[row, row] = k1
                A[row, row+1] = -1
            else:
                # In bar2 region
                A[row, row-1] = -1
                A[row, row] = k2
                A[row, row+1] = -1
    return A

def fem_array(bar1:'Bar', bar2:'Bar', num_sections: int=6, x_bar=2/np.pi, l=1):
    '''num_sections = total sections to cut the bar. Note that this is NOT evenly distributed over the entire
    bar. Half of the sections are left of x_bar, half are to the right.
    Returns the A matrix of the FEM equation Ax = B'''
    dx1 = x_bar / (num_sections / 2)
    dx2 =  (l - x_bar) / (num_sections / 2)
    k1 = (2/dx1 + 4*bar1.alpha**2*dx1/6) / (1/dx1 - bar1.alpha**2*dx1/6)
    k2 = (2/dx2 + 4*bar2.alpha**2*dx2/6) / (1/dx2 - bar2.alpha**2*dx2/6)

    beta1 = (bar1.k * bar1.area) / (dx1)
    beta2 = (bar2.k * bar2.area) / (dx2)

    rho = beta1 + beta2 + (bar1.k * bar1.area * dx1 * bar1.alpha**2) / 2 + (bar2.k * bar2.area * dx2 * bar2.alpha**2) / 2
    # Making the matrix
    # There are num_sections - 1 rows and columns
    # row 0, row n - 2, and row num_sections/2 - 1 are unique
    # the rest are -1, k, -1 where k changes from k1 to k2 at num_sections/2

    # Initialize matrix A with zeros
    A = np.zeros((num_sections - 1, num_sections - 1))
    
    # Fill the matrix A
    for row in range(num_sections - 1):
        if row == 0:
            # First row for boundary at x = 0 (for example Dirichlet or Neumann condition)
            A[row, row] = k1
            A[row, row+1] = -1
        elif row == num_sections // 2 - 1:
            # Special row at the interface between bar1 and bar2
            A[row, row-1] = -beta1
            A[row, row] = rho  # This is the rho value you computed
            A[row, row+1] = -beta2
        elif row == num_sections - 2:
            # Last row for boundary at x = l (again assuming Dirichlet or Neumann condition)
            A[row, row-1] = -1
            A[row, row] = k2
        else:
            # Interior rows for bar1 and bar2 (uniform grid)
            if row < num_sections // 2 - 1:
                # In bar1 region
                A[row, row-1] = -1
                A[row, row] = k1
                A[row, row+1] = -1
            else:
                # In bar2 region
                A[row, row-1] = -1
                A[row, row] = k2
                A[row, row+1] = -1
    return A

def solve(bar1:'Bar', bar2:'Bar', num_sections: int=6, x_bar=2/np.pi, l=1, method="FDM"):
    '''Solves using FDM or FEM as specified. Returns (temps, x_values)'''
    if method == "FDM":
        A = fdm_array(bar1,
                      bar2,
                      num_sections=num_sections,
                      x_bar=x_bar,
                      l=l)
    elif method == "FEM":
        A = fem_array(bar1,
                      bar2,
                      num_sections=num_sections,
                      x_bar=x_bar,
                      l=l)
    
    B = np.zeros((num_sections - 1))
    B[-1] = 100

    # Add boundary conditions
    interior_temps = np.linalg.solve(A, B)
    all_temps = np.array([0])
    all_temps = np.append(all_temps, interior_temps)
    all_temps = np.append(all_temps, 100)

    # Generate the x_values to return along with the temps
    dx2 =  (l - x_bar) / (num_sections / 2)
    x_values = np.linspace(0, x_bar, num_sections // 2 + 1) # [0, x_bar]
    x_values = np.append(x_values, np.linspace(x_bar+dx2, l, num_sections // 2))
    return all_temps, x_values


if __name__ == "__main__":
    bar1 = Bar(
        radius=0.1,
        k=0.5,
        h=0.25
    )
    bar2 = Bar(
        radius=0.1,
        k=2,
        h=0.25
    )

    set_k_ratio(0.5, bar1, bar2)

    x_bar = 0.5

    x_values = np.linspace(0, 1, 1000)
    temp_values = get_analytical_datapoints(x_points=x_values,
                                            bar1=bar1,
                                            bar2=bar2,
                                            x_bar=x_bar)


    # fdm tests
    l = 1
    num_sections = 4
    dx1 = x_bar / (num_sections / 2)
    dx2 =  (l - x_bar) / (num_sections / 2)
    fdm_temps, x_fdm = solve(bar1=bar1,
                    bar2=bar2,
                    num_sections=num_sections,
                    x_bar=x_bar,
                    l=1,
                    method="FDM")
    fem_temps, x_fem = solve(bar1=bar1,
                    bar2=bar2,
                    num_sections=num_sections,
                    x_bar=x_bar,
                    l=1,
                    method="FEM")

    # Plot the data
    plt.plot(x_values, temp_values, label='Temperature Distribution')
    plt.axvline(x=x_bar, color='r', linestyle='--', label='x_bar')
    plt.plot(x_fdm, fdm_temps, 'o-', label='Temperature Distribution (FDM)', color='g')
    plt.plot(x_fem, fem_temps, 'o-', label='Temperature Distribution (FEM)', color='r')
    plt.xlabel('x')
    plt.ylabel('Temperature')
    plt.title('Temperature Distribution Along the Bar')
    plt.legend()
    plt.grid(True)

    # Show plot
    plt.show()