import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

sns.set_theme()
plt.rcParams['text.usetex'] = True


def plot_transformations(save: bool=False) -> None:
    # Load data
    tridiag = pd.read_csv("../latex/output/transform_tridiag.csv", header=0)
    dense = pd.read_csv("../latex/output/transform_dense.csv", header=0)

    fig, ax = plt.subplots()
    ax.loglog(dense['N'], dense['T'], '--', label='Dense')
    ax.loglog(tridiag['N'], tridiag['T'], label='Tridiagonal')
    ax.set_xlabel('N')
    ax.set_ylabel('Similarity transformations')
    ax.legend()

    # Save to file
    if save is True:
        fig.savefig("../latex/images/transform.pdf")


def plot_eigenvectors(N: int, save: bool=False) -> None:
    # Load data based on matrix size
    path = f"../latex/output/eigenvector_{N}.csv"
    eigvec = pd.read_csv(path, header=0)

    fig, ax = plt.subplots()
    ax.plot(eigvec['x'], eigvec['Vector 1'], label=r'$\vec{v}_{1}$')
    ax.plot(eigvec['x'], eigvec['Vector 2'], label=r'$\vec{v}_{2}$')
    ax.plot(eigvec['x'], eigvec['Vector 3'], label=r'$\vec{v}_{3}$')
    ax.plot(eigvec['x'], eigvec['Analytic 1'], '--', label=r'$\vec{v}^{(1)}$')
    ax.plot(eigvec['x'], eigvec['Analytic 2'], '--', label=r'$\vec{v}^{(2)}$')
    ax.plot(eigvec['x'], eigvec['Analytic 3'], '--', label=r'$\vec{v}^{(3)}$')
    ax.set_xlabel(r'Element $\hat{x}_{i}$')
    ax.set_ylabel(r'Element $v_{i}$')
    ax.legend(loc='upper left')

    # Save to file
    if save is True:
        fig.savefig(f"../latex/images/eigenvector_{N}.pdf")

     
if __name__ == '__main__':
    plot_transformations(True)
    plot_eigenvectors(10, True)
    plot_eigenvectors(100, True)