import os

import numpy as np
import scipy.io


class MatrixGenerator:
    @staticmethod
    def generate_spd_matrix(n: int) -> np.ndarray:
        A = np.random.rand(n, n)
        return np.dot(A, A.T) + n * MatrixGenerator.generate_identity_matrix(
            n
        )  # Adding n*I ensures positive definiteness

    @staticmethod
    def generate_identity_matrix(size):
        return np.eye(size)

    @staticmethod
    def generate_alternating_vector(size):
        return np.tile([1, 2], int(np.ceil(size / 2)))[:size]

    @staticmethod
    def get_matrix_from_file(file_path, problem):
        mat_contents = scipy.io.loadmat(file_path)
        problem_record = mat_contents["Problem"][0][0]
        A = problem_record[problem]
        if scipy.sparse.issparse(A):
            A_dense = A.todense()
        else:
            A_dense = A
        return np.array(A_dense)

    @staticmethod
    def generate_matrix_and_vector(type, size=None):
        if type == "spd":
            if size is None:
                raise ValueError("Size must be provided for SPD matrix generation.")
            try:
                matrix, vector, lambda_min, lambda_max = MatrixGenerator.load_from_file(
                    "spd_" + str(size) + ".npz"
                )
            except FileNotFoundError:
                matrix = MatrixGenerator.generate_spd_matrix(size)
                vector = np.random.uniform(-1, 1, size)
                lambda_min, lambda_max = MatrixGenerator.calculate_eigenvalues(matrix)
                MatrixGenerator.save_to_file(
                    matrix, vector, lambda_min, lambda_max, "spd_" + str(size) + ".npz"
                )
        elif type == "nemeth12":
            try:
                matrix, vector, lambda_min, lambda_max = MatrixGenerator.load_from_file(
                    "nemeth12.npz"
                )
            except FileNotFoundError:
                matrix = -1 * MatrixGenerator.get_matrix_from_file("nemeth12.mat", 1)
                size = matrix.shape[0]
                vector = MatrixGenerator.generate_alternating_vector(size)
                lambda_min, lambda_max = MatrixGenerator.calculate_eigenvalues(matrix)
                MatrixGenerator.save_to_file(
                    matrix, vector, lambda_min, lambda_max, "nemeth12.npz"
                )
        elif type == "poli3":
            try:
                matrix, vector, lambda_min, lambda_max = MatrixGenerator.load_from_file(
                    "poli3.npz"
                )
            except FileNotFoundError:
                matrix = MatrixGenerator.get_matrix_from_file("poli3.mat", 2)
                size = matrix.shape[0]
                vector = MatrixGenerator.generate_alternating_vector(size)
                lambda_min, lambda_max = MatrixGenerator.calculate_eigenvalues(matrix)
                MatrixGenerator.save_to_file(
                    matrix, vector, lambda_min, lambda_max, "poli3.npz"
                )
        else:
            raise ValueError(
                "Invalid type specified. Choose 'spd', 'nemeth12', or 'poli3'."
            )

        return matrix, vector, lambda_min, lambda_max

    @staticmethod
    def calculate_eigenvalues(A):
        eigenvalues = np.linalg.eigvals(A)
        lambda_min = np.min(eigenvalues)
        lambda_max = np.max(eigenvalues)
        return lambda_min, lambda_max

    @staticmethod
    def save_to_file(matrix, vector, lambda_min, lambda_max, file_path) -> None:
        np.savez(
            file_path,
            matrix=matrix,
            vector=vector,
            lambda_min=lambda_min,
            lambda_max=lambda_max,
        )

    def load_from_file(file_path):
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"The file {file_path} does not exist.")
        data = np.load(file_path)
        matrix = data["matrix"]
        vector = data["vector"]
        lambda_min = data["lambda_min"]
        lambda_max = data["lambda_max"]
        return matrix, vector, lambda_min, lambda_max