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='Vector 1')
    ax.plot(eigvec['x'], eigvec['Vector 2'], label='Vector 2')
    ax.plot(eigvec['x'], eigvec['Vector 3'], label='Vector 3')
    ax.set_xlabel(r'Element $\hat{x}_{i}$')
    ax.set_ylabel(r'Value of element $v_{i}$')
    ax.legend()

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

     


if __name__ == '__main__':
    plot_transformations(True)
    plot_eigenvectors(6, True)
    plot_eigenvectors(100, True)
    # plt.show()