# 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, order=2):
    '''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)

    if order == 2:
        w1 = dx1
        w2 = dx2

        epsilon1 = 1 + dx1**2 * bar1.alpha**2 / 2
        epsilon2 = 1 + dx2**2 * bar2.alpha**2 / 2

        beta1 = bar1.k * bar1.area / w1
        beta2 = bar2.k * bar2.area / w2

        k1 = 2*beta1*epsilon1
        k2 = 2*beta2*epsilon2

        rho = beta1*epsilon1 + beta2*epsilon2
    elif order == 4:
        w1 = dx1 + dx1**3 * bar1.alpha**2 / 6
        w2 = dx2 + dx2**3 * bar2.alpha**2 / 6

        epsilon1 = 1 + dx1**2 * bar1.alpha**2 / 2 + dx1**4 * bar1.alpha**4 / 24
        epsilon2 = 1 + dx2**2 * bar2.alpha**2 / 2 + dx2**4 * bar2.alpha**4 / 24

        beta1 = bar1.k * bar1.area / w1
        beta2 = bar2.k * bar2.area / w2

        k1 = 2*beta1*epsilon1
        k2 = 2*beta2*epsilon2

        rho = beta1*epsilon1 + beta2*epsilon2

    # 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] = -beta1
        elif row == num_sections // 2 - 1:
            # Special row at the interface between bar1 and bar2
            A[row, row-1] = -beta1
            A[row, row] = rho
            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] = -beta2
            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] = -beta1
                A[row, row] = k1
                A[row, row+1] = -beta1
            else:
                # In bar2 region
                A[row, row-1] = -beta2
                A[row, row] = k2
                A[row, row+1] = -beta2
    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 get_fem_p2_k(row, col, bar:'Bar', dx):
    if (row == 1 and col == 1) or (row == 2 and col == 2):
        return bar.k * bar.area / dx + bar.h * bar.p * dx / 3.0
    elif row == 3 and col == 3:
        return 1 * bar.k * bar.area / (3*dx) + bar.h * bar.p * dx / 30
    elif (row == 1 and col == 2) or (row == 2 and col == 1):
        return -1 * bar.k * bar.area / dx + bar.h * bar.p * dx / 6
    elif (row == 1 and col == 3) or (row == 3 and col == 1) or (row == 2 and col == 3) or (row == 3 and col == 2):
        return bar.h * bar.p * dx / 12
    
def get_all_fem_p2_k(bar:'Bar', dx):
    K = np.zeros((3, 3))
    for row in range(3):
        for col in range(3):
            K[row, col] = get_fem_p2_k(row+1, col+1, bar, dx)
    return K

def fem_p_2_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)

    K_1 = get_all_fem_p2_k(bar1, dx1)
    K_2 = get_all_fem_p2_k(bar2, dx2)
    
    # Initialize matrix A with zeros
    rows = 2 * num_sections - 1
    cols = 2 * num_sections - 1
    A = np.zeros((rows, cols))

    omega_col_index = num_sections - 1
    num_omega = num_sections
    for i in range(num_sections):
        row1 = i*2
        row2 = row1 + 1
        if i == 0:
            A[row1][0] = K_1[2-1, 2-1] + K_1[1-1, 1-1]
            A[row1][1] = K_1[1-1, 2-1]
            A[row1][omega_col_index] = K_1[2-1, 3-1]
            A[row1][omega_col_index+1] = K_1[1-1, 3-1]

            A[row2][0] = K_1[3-1, 2-1]
            A[row2][omega_col_index] = K_1[3-1, 3-1]
        elif i == num_sections - 2: # Second last pair of rows
            A[row1][i-1] =                  K_2[2-1, 1-1]
            A[row1][i] = K_2[2-1, 2-1] +    K_2[1-1, 1-1]
            A[row1][omega_col_index+i] =    K_2[2-1, 3-1]
            A[row1][omega_col_index+i+1] =  K_2[1-1, 3-1]

            A[row2][i-1] = K_2[3-1, 1-1]
            A[row2][i] = K_2[3-1, 2-1]
            A[row2][omega_col_index+i] = K_2[3-1, 3-1]
        elif i == num_sections - 1: # Last pair of rows, use only omega equation
            A[row1][omega_col_index-1] = K_2[3-1, 1-1]
            A[row1][-1] = K_2[3-1, 3-1]
        elif i == num_sections // 2 - 1:  # Interface element
            A[row1][i-1] = K_1[2-1, 1-1]
            A[row1][i] = K_1[2-1, 2-1] + K_2[1-1, 1-1]
            A[row1][i+1] = K_2[1-1, 2-1]
            A[row1][omega_col_index + i] = K_1[2-1, 3-1]
            A[row1][omega_col_index + i + 1] = K_2[1-1, 3-1]
            
            A[row2][i-1] = K_1[3-1, 1-1]
            A[row2][i] = K_1[3-1, 2-1]
            A[row2][omega_col_index + i] = K_1[3-1, 3-1]
        else:
            if i < num_sections // 2: # We are in the first bar
                A[row1][i-1] = K_1[2-1, 1-1]
                A[row1][i] = K_1[2-1, 2-1] + K_1[1-1, 1-1]
                A[row1][i+1] = K_1[1-1, 2-1]
                A[row1][omega_col_index + i] = K_1[2-1, 3-1]
                A[row1][omega_col_index + i + 1] = K_1[1-1, 3-1]

                A[row2][i-1] = K_1[3-1, 1-1]
                A[row2][i] = K_1[3-1, 2-1]
                A[row2][omega_col_index + i] = K_1[3-1, 3-1]
            else:   # We are in the second bar
                A[row1][i-1] = K_2[2-1, 1-1]
                A[row1][i] = K_2[2-1, 2-1] + K_2[1-1, 1-1]
                A[row1][i+1] = K_2[1-1, 2-1]
                A[row1][omega_col_index + i] = K_2[2-1, 3-1]
                A[row1][omega_col_index + i + 1] = K_2[1-1, 3-1]

                A[row2][i-1] = K_2[3-1, 1-1]
                A[row2][i] = K_2[3-1, 2-1]
                A[row2][omega_col_index + i] = K_2[3-1, 3-1]
    return A
        


def solve(bar1:'Bar', bar2:'Bar', num_sections: int=6, x_bar=2/np.pi, l=1, method="FDM", order=2):
    '''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,
                      order=order)
        # Generate the B vector
        if order == 2:
            dx2 =  (l - x_bar) / (num_sections / 2)
            w2 = dx2
            beta2 = bar2.k * bar2.area / w2
        elif order == 4:
            dx2 =  (l - x_bar) / (num_sections / 2)
            w2 = dx2 + dx2**3 * bar2.alpha**2 / 6
            beta2 = bar2.k * bar2.area / w2
        
        B = np.zeros((num_sections - 1))
        B[-1] = beta2*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))
        
    elif method == "FEM":
        if order == 2:
            A = fem_array(bar1,
                      bar2,
                      num_sections=num_sections,
                      x_bar=x_bar,
                      l=l)
            # Generate the B vector
            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))
        elif order == 4:
            A = fem_p_2_array(bar1,
                      bar2,
                      num_sections=num_sections,
                      x_bar=x_bar,
                      l=l)

            # Generate the B vector
            dx2 =  (l - x_bar) / (num_sections / 2)
            B = np.zeros((num_sections*2-1))
            K_2 = get_all_fem_p2_k(bar2, dx2)
            B[0] = 0
            B[1] = 0
            B[-3] = -1 * K_2[1-1, 2-1] * 100
            B[-2] = 0
            B[-1] = -1 * K_2[3-1, 2-1] * 100

            # Add boundary conditions
            solution = np.linalg.solve(A, B)
            all_temps = np.array([0])
            all_temps = np.append(all_temps, solution[0:num_sections-1])
            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=0.5,
        h=0.25
    )

    #set_k_ratio(8, bar1, bar2)
    

    x_bar = 0.7

    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 = 16
    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="FEM",
                    order=2)
    fem_temps, x_fem = solve(bar1=bar1,
                    bar2=bar2,
                    num_sections=num_sections,
                    x_bar=x_bar,
                    l=1,
                    method="FEM",
                    order=4)
    
    print(f"A: {bar1.area}")
    print(f"k: {bar1.k}")
    print(f"h: {bar1.h}")
    print(f"P: {bar1.p}")
    print(f"dx: {dx1}")


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