import matplotlib.pyplot as plt
import math
import numpy as np
import time
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize
from scipy.sparse.linalg import cg
from scipy.sparse.linalg import gmres

def postpro(results, n, m):
    plt.close('all')
    results = np.asarray(results)

    if results.size != n * m:
        raise ValueError("results size must be n * m")

    field = np.flipud(results.reshape((n, m)))


    x = np.linspace(0, domain_length_x, m)
    y = np.linspace(0, domain_length_y, n)
    X, Y = np.meshgrid(x, y)

    fig, ax = plt.subplots(figsize=(12, 6))

    cmap = plt.cm.turbo.copy()
    cmap.set_bad(color="white")  # masked values → white

    im = ax.pcolormesh(
        X, Y, field,
        cmap=cmap,
        shading="nearest",
        norm=Normalize(results.min(), results.max())
    )

    ax.set_xlabel("domain_length")
    ax.set_ylabel("domain_height")
    ax.set_title("Temperature Plot")

    cbar = plt.colorbar(im, ax=ax, pad=0.02)
    cbar.set_label("Temperature (K)")
    cbar.ax.locator_params(nbins=8)
    cbar.update_ticks()

    ax.set_aspect("auto")
    plt.tight_layout()
    plt.show()


def one2d(number, n, m):
    x = math.floor((number-1)/m)
    y = number - x * m - 1
    return [x,y]

def two1d(x,y,n,m):
    return x*m+y + 1

def sin_deg(degrees):
    return math.sin(math.radians(degrees))

def GaussSeidl(A,x,b,max_iter):
    n = len(A)
    A = np.array(A, dtype=float)
    x = np.array(x0, dtype=float)
    b = np.array(b, dtype=float)
    init = 0


    for k in range(max_iter):
        start = time.perf_counter()
        for i in range(n):
            s = 0.0
            for j in range(n):
                if(i!=j):
                    s += x[j] * A[i][j]
            x[i] = (b[i]-s) / A[i,i]

        #Residual calculation
        r = np.zeros(len(b))
        for l in range(n):
            cr = 0
            for i in range(n):
                cr += x[i] * A[l][i]
            r[l] = cr-b[l]
        print(r)
        res = 0
        for el in r:
            res += el**2
        res = res/len(r)
        res = math.sqrt(res)

        if k == 0:
            init = res
        res = res/init
        n2 = round(len(A)**0.5)
        #postpro(r/init,n2,n2)
        #postpro(x,n2,n2)
        end = time.perf_counter()
        t = (end - start)
        print(f'{k+1}. Residual: {res} Time: {t:.6e}')


    return x



# =========================
# mesh settings
# =========================
base_size = 0.01  # m
domain_length_x = 0.3
domain_length_y = 0.3
n = round(domain_length_x / base_size)
m = round(domain_length_y / base_size)
print("Cell count: " + str(n*m))

# =========================
# boundary conditions
# =========================
boundary = np.zeros(((n+2),(m+2)))
velocity = 1
velocity_angle = 0
avg = (273*3 + 353)/4
x0 = [avg for i in range(m*n)]


# top temp
top_temp = 273
for i in range(n*m):
    if i < m:
        boundary[0][i % m + 1] = top_temp

# bottom temp
bottom_temp = 273
for i in range(n*m):
    if i > m * (n-1) - 1:
        boundary[n+1][i % m + 1] = bottom_temp

# left temp
left_temp = 353
for i in range(n*m):
    if (i-1) % m == 0:
        boundary[math.floor((i)/m) + 1][0] = left_temp

# right temp
right_temp = 273
for i in range(n*m):
    if (i) % m == 0:
        boundary[math.floor((i)/m) + 1][m+1] = right_temp

"""
for i in range(10):
    boundary[27+i][0] = 400
print(boundary)
"""

# =========================
# physics values
# =========================
thermal_conductivity = 960  # W/m/K
area = 1  # m^2
density = 7800

# =========================
# discretisation
# =========================
matrix = np.zeros((n*m, n*m))
nodes = np.zeros(n*m)
init = np.zeros(n*m)

for i in range(0, m*n):
    x = one2d(i+1, n, m)[0] + 1
    y = one2d(i+1, n, m)[1] + 1
    a_p = 0

    # left cell
    if boundary[x][y-1] == 0:
        a_p += thermal_conductivity * area / base_size + 0.5 * sin_deg(velocity_angle+90) * velocity * density
        matrix[i][i - 1] = -(thermal_conductivity * area / base_size  + sin_deg(velocity_angle+90) * velocity * density)
    else:
        a_p += 2 * thermal_conductivity * area / base_size  + 1 * sin_deg(velocity_angle+90) * velocity * density
        init[i] += boundary[x][y-1] * (2 * thermal_conductivity * area / base_size  + sin_deg(velocity_angle+90) * velocity * density)

    # top cell
    if boundary[x-1][y] == 0:
        a_p += thermal_conductivity * area / base_size  + 0.5 * sin_deg(-velocity_angle) * velocity * density
        matrix[i][i - n] = -(thermal_conductivity * area / base_size  + sin_deg(-velocity_angle) * velocity * density)
    else:
        a_p += 2 * thermal_conductivity * area / base_size  + 1 * sin_deg(-velocity_angle) * velocity * density
        init[i] += boundary[x-1][y] * (2 * thermal_conductivity * area / base_size  + sin_deg(-velocity_angle) * velocity * density)

    # right cell
    if boundary[x][y+1] == 0:
        a_p += thermal_conductivity * area / base_size  + 0.5 * sin_deg(velocity_angle-90) * velocity * density
        matrix[i][i + 1] = -(thermal_conductivity * area / base_size  + sin_deg(velocity_angle-90) * velocity * density)
    else:
        a_p += 2 * thermal_conductivity * area / base_size + 1 * sin_deg(velocity_angle-90) * velocity * density
        init[i] += boundary[x][y+1] * (2 * thermal_conductivity * area / base_size  + sin_deg(velocity_angle-90) * velocity * density)

    # bottom cell
    if boundary[x+1][y] == 0:
        a_p += thermal_conductivity * area / base_size + 0.5 * sin_deg((velocity_angle)) * velocity * density
        matrix[i][i + n] = -(thermal_conductivity * area / base_size + sin_deg((velocity_angle)) * velocity * density)
    else:
        a_p += 2 * thermal_conductivity * area / base_size + 1 * sin_deg((velocity_angle)) * velocity * density
        init[i] += boundary[x+1][y] * (2 * thermal_conductivity * area / base_size  + sin_deg((velocity_angle)) * velocity * density)

    matrix[i][i] = a_p



# =========================
# Solver
# =========================

#nodes = np.linalg.solve(matrix, init)
nodes = GaussSeidl(matrix,x0,init,100)

# =========================
# Results
# =========================
print(matrix)
print(init)
print(nodes)


# =========================
# Post Processing
# =========================


postpro(nodes, n, m)