import subprocess
import tyro
from dataclasses import dataclass
from pathlib import Path
from nerfstudio.utils.eval_utils import eval_setup
from typing import Optional
import pandas as pd

def verify_pdflatex_installed():
    """Check if pdflatex is installed."""
    try:
        subprocess.run(
            ["pdflatex", "--version"],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        print("pdflatex is installed.")
        return True
    except subprocess.CalledProcessError:
        print("pdflatex is not installed. Please install it and try again.")
        return False


def renders_to_latex(result_folder):
    # Convert result_folder to Path object if it's a string
    result_folder = Path(result_folder)

    # Initialize the groups with specified labels
    groups = {
        'depth': {
            'folders': ['gt_depth', 'depth_renders', 'mesh_renders'],
            'headers': ['gt depth', 'depth render', 'mesh render']
        },
        'normal': {
            'folders': ['gt_surface', 'surface_renders', 'normal_renders'],
            'headers': ['gt surface', 'surface render', 'normal render']
        },
        'rgb': {
            'folders': ['gt_images', 'rgb_renders'],
            'headers': ['gt image', 'rgb render']
        }
    }

    reference_folder = next((result_folder / folder for folder in ['rgb_renders', 'gt_images'] if (result_folder / folder).exists()), None)

    # Use the 'rgb_renders' folder to generate the image name list
    print(f"Using '{reference_folder}' to generate image name list.\n")

    # List all image files in the reference folder
    image_extensions = {'.png', '.jpg', '.jpeg'}
    image_files = [f for f in reference_folder.iterdir() if f.suffix.lower() in image_extensions]

    # Extract just the image filenames
    image_names = [f.name for f in image_files]
    image_names.sort()
    print(f"Image names: {image_names}\n")
    # Generate LaTeX code for each group
    latex_code = ""

    for image_name in image_names:
        image_description = image_name.replace("_", " ")
        latex_code += f"\\section{{Image: {image_description}}}\n\n"
        
        for group_name, group_info in groups.items():
            folders = group_info['folders']
            headers = group_info['headers']
            
            num_columns = len(headers)
            latex_code += f"\\begin{{figure}}[H]\n\\centering\n"
            
            for i, (folder, header) in enumerate(zip(folders, headers)):
                image_path = Path(folder).joinpath(image_name)
                latex_code += f"\\begin{{subfigure}}[b]{{0.3\\textwidth}}\n"
                latex_code += f"    \\centering\n"
                latex_code += f"    \\includegraphics[width=\\textwidth, height=\\textwidth, keepaspectratio]{{{image_path.as_posix()}}}\n"
                latex_code += f"    \\caption{{{header}}}\n"
                latex_code += f"\\end{{subfigure}}\n"
                if i < num_columns - 1:
                    latex_code += "\\hfill\n"
            
            latex_code += f"\\caption{{{group_name.capitalize()} group for {image_description}}}\n"
            latex_code += "\\end{figure}\n\n"
        
        latex_code += "\\clearpage\n\n"  # Add a page break after each image set

    return latex_code


def compile_latex(tex_file):
    """Compile the LaTeX file to generate a PDF."""
    try:
        parent_dir = tex_file.parent
        subprocess.run(
            ["pdflatex", "-interaction=nonstopmode", tex_file],
            check=True,
            cwd=parent_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        extensions = [".aux", ".log", ".out", ".toc"]

        compressed_file = parent_dir.joinpath(f"{tex_file.stem}_compressed.pdf")
        pdf_file = parent_dir.joinpath(f"{tex_file.stem}.pdf")
        
        # Use Ghostscript to compress the PDF
        subprocess.run([
            "gs", "-sDEVICE=pdfwrite", "-dCompatibilityLevel=1.4",
            "-dPDFSETTINGS=/ebook", "-dNOPAUSE", "-dQUIET", "-dBATCH",
            f"-sOutputFile={compressed_file}", pdf_file
        ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        for ext in extensions:
            artifact = tex_file.with_suffix(ext)
            if artifact.exists():
                artifact.unlink()

    except subprocess.CalledProcessError as e:
        print(f"Error during LaTeX compilation: {e.stderr.decode()}")
        return False
    return True


def config_to_latex(config):
    """Convert the configuration dictionary to a LaTeX block with JSON-style code."""
    config_str = str(config)
    latex_block = r"\section*{Configuration}" + "\n"
    latex_block += r"\begin{verbatim}" + "\n"
    latex_block += config_str + "\n"
    latex_block += r"\end{verbatim}" + "\n"
    return latex_block      


def dict_to_latex_table(data):
    # Table header
    table = r"\section*{Metrics}" + "\n"
    table += r"\begin{table}[ht]" + "\n"
    table += r"\centering" + "\n"
    table += r"\begin{tabular}{|l|c|c|}" + "\n"
    table += r"\hline" + "\n"
    table += r"\textbf{Metric} & \textbf{Mean} & \textbf{Std} \\" + "\n"
    table += r"\hline" + "\n"
    
    # Loop through the dictionary to add rows
    processed_metrics = set()  # Keep track of metrics already processed
    for key in data:
        if '_std' in key:
            metric = key.replace('_std', '')
            if metric not in processed_metrics:
                table += f"{metric.replace('_', ' ').upper()} & {data[metric]:.4f} & {data[key]:.4f} \\\\" + "\n"
                processed_metrics.add(metric)
        elif key not in processed_metrics and f"{key}_std" not in data:
            table += f"{key.replace('_', ' ').upper()} & {data[key]} & \\\\" + "\n"
            processed_metrics.add(key)

    # End of table
    table += r"\hline" + "\n"
    table += r"\end{tabular}" + "\n"
    table += r"\caption{Metrics and their corresponding mean and standard deviation.}" + "\n"
    table += r"\label{tab:metrics}" + "\n"
    table += r"\end{table}" + "\n"

    return table

def profiler_csv_to_latex(profile_csv_path):
    """
    Convert the profiler CSV file to a LaTeX table.
    
    Load the CSV into pandas.
    Convert to LaTeX.
    Add the header "Profiler Metrics".
    Return the LaTeX code.
    """
    if profile_csv_path is None:
        return ""
    latex_table = r"\section*{Profiler Metrics (FPS)}" + "\n"
    df = pd.read_csv(profile_csv_path)
    pd.options.display.float_format = '{:.2f}'.format
    latex_table += df.to_latex(index=False, float_format='%.2f')
    return latex_table


def write_latex(file_path, config, metrics, profile_csv_path):
    """Write LaTeX content to a .tex file."""
    latex_content = r"\documentclass{article}" + "\n"
    latex_content += r"\usepackage{verbatim}" + "\n"
    latex_content += r"""
\usepackage{graphicx}
\usepackage{float}
\usepackage{geometry}
\usepackage{booktabs}
\usepackage{subcaption}
\usepackage{caption}
"""
    latex_content += r"\begin{document}" + "\n"
    latex_content += dict_to_latex_table(metrics)
    latex_content += profiler_csv_to_latex(profile_csv_path)
    latex_content += renders_to_latex(file_path.parent)
    latex_content += config_to_latex(config)
    latex_content += r"\end{document}" + "\n"

    with open(file_path, 'w') as file:
        file.write(latex_content)


def update_config_callback(config):
    config.pipeline.datamanager.camera_res_scale_factor = 0.50
    if config.pipeline.model.camera_optimizer.mode == "off":
        config.pipeline.datamanager.dataparser.train_split_fraction = 0.9
    return config


@dataclass
class ReportGenerator:
    run_dir: Path
    """Directory pointing to the last run."""

    report_dir: Path = Path().home().joinpath("reports")
    """Directory where reports are generated."""

    output_dir: Optional[Path] = None
    """Directory where reports are generated."""


    def main(self):
        report_dir = self.report_dir
        if not report_dir.exists():
            config.report_dir.mkdir(exist_ok=True)

        config_path = self.run_dir.joinpath("config.yml")
        assert config_path.exists(), f"This run is invalid: {config_path} doesn't exist."

        config, pipeline, checkpoint_path, _ = eval_setup(config_path, update_config_callback=update_config_callback)

        if len(pipeline.datamanager.fixed_indices_eval_dataloader) == 0:
            metrics = {}
        else:
            metrics = pipeline.get_average_eval_image_metrics(
                output_path=None,
                get_std=True
            )

        output_dir = self.output_dir
        if output_dir is None:
            output_dir = report_dir.joinpath(config.experiment_name)
            output_dir.mkdir(parents=True, exist_ok=True)

        file_name = output_dir.joinpath(config.timestamp + ".tex")
        
        prorfile_csv_path = None
        if pipeline.model.config.enable_profiling:
            profile_csv_path = output_dir.joinpath(f"profile_outputs.csv")
        write_latex(file_name , config, metrics, profile_csv_path)

        compile_latex(file_name)


if __name__ == "__main__":
    if not verify_pdflatex_installed():
        raise FileNotFoundError("pdflatex doesn't exists. Install tex-live.")

    config = tyro.cli(ReportGenerator).main()
