import argparse
from pathlib import Path
from matplotlib import pyplot as plt
from stereograph.plots import combine_energy_resolutions, combine_mean_biases, combine_angular_resolutions, combine_auc
import matplotlib
matplotlib.use('Agg')


def main(file_paths, outfile, energy=False, direction=False, classification=False):
    # fig, axes = plt.subplots(2, 2, figsize=(20, 12))

    if energy:
        fig, axes = plt.subplots(figsize=(18, 6), nrows=1, ncols=2)
        ax = axes[0]
        ax = combine_energy_resolutions(file_paths, ax=ax, label="graphs")
        ax2 = axes[1]
        combine_mean_biases(file_paths, ax=ax2, ls="--")
    
    if direction:
        plt.figure(figsize=(12, 6))
        ax = None
        combine_angular_resolutions(file_paths, ax=ax, label="graphs")

    if classification:
        plt.figure(figsize=(12, 6))
        ax = None
        combine_auc(file_paths, ax=ax)
    
    plt.tight_layout()
    plt.savefig(outfile, dpi=200)
    

def parse_arguments():
    parser = argparse.ArgumentParser(description="Make performance plots from DL2 data files.\n" \
        "Example usage: python script_plot_performances.py -f dl2.h5 -e -d -c -o output.png")

    parser.add_argument("--filenames", "-f", type=Path, nargs='+', help="Paths to the DL2 files with prediction values")
    parser.add_argument("--energy", "-e", action="store_true", help="Make energy resolution plot if passed")
    parser.add_argument("--direction", "-d", action="store_true", help="Make angular resolution plot if passed")
    parser.add_argument("--classification", "-c", action="store_true", help="Make AUC plot if passed")
    parser.add_argument("--output", "-o", type=Path, help="Path to the output plot file", default="resolution.png")
    parser.add_argument("--overwrite", action="store_true", help="Overwrite the output file if it exists")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_arguments()
    
    outfile = args.output
    outfile.parent.mkdir(parents=True, exist_ok=True)

    if not args.overwrite and outfile.exists():
        raise FileExistsError(
            f"Output {outfile} already exists. Use --overwrite to overwrite it."
        )

    main(args.filenames, outfile, args.energy, args.direction, args.classification)