#!/usr/bin/env python3

# this script has been tested with
# - python 3.11.5
# required packages
# - pandas
# - openpyxl

import pandas as pd
import sys
import os

import argparse


print("Saving files in:", os.getcwd())

QP_COL = 'D'
RATE_COL = 'E'
PSNR_Y_COL ='F'
PSNR_U_COL ='G'
PSNR_V_COL ='H'
NAME_COL1 = 'A'
NAME_COL2 = 'B'

def get_col_index(col_letter):
    return ord(col_letter.upper())-ord('A')


def extract_all(filepath, csv_format, max_rows=120):
    xls = pd.ExcelFile(filepath)

    if csv_format == "normal":
        for sheet_name in xls.sheet_names:
            if sheet_name in [
                "Summary",
                "Plot-trend",
                "Plot - PSNR",
                "RA-10",
                "LDB-10",
                "Reference",
                "Test",
                "PerSequence",
            ]:
                print(f"Skipping {sheet_name}")
                continue

            df = pd.read_excel(xls, sheet_name=sheet_name)
            clean_name = sheet_name.strip()
            csv_file = f"{clean_name}.csv"
            df.to_csv(csv_file, index=False, sep=",")
            print(f"Exported '{sheet_name}' to '{csv_file}'")

    if csv_format == "rdplot":
        for sheet_name in xls.sheet_names:
            if sheet_name in [
                "Summary",
                "Plot-trend",
                "Plot - PSNR",
                "RA-10",
                "LDB-10",
                "Reference",
                "Test",
                "PerSequence",
            ]:
                print(f"Skipping {sheet_name}")
                continue

            all_data = []

            df = pd.read_excel(xls, sheet_name=sheet_name, header=None, nrows=max_rows)

            for i, row in df.iterrows():
                try:
                    qp = row[get_col_index(QP_COL)]
                    rate = row[get_col_index(RATE_COL)]
                    psnr_y = row[get_col_index(PSNR_Y_COL)]
                    psnr_u = row[get_col_index(PSNR_U_COL)]
                    psnr_v = row[get_col_index(PSNR_V_COL)]

                    name_part1 = row[get_col_index(NAME_COL1)]
                    name_part2 = row[get_col_index(NAME_COL2)]
                    if pd.isna(name_part1) or pd.isna(name_part2):
                        continue
                    name = f"{name_part1} {name_part2}"

                    if pd.isna(rate) or pd.isna(psnr_y):
                        continue

                    if pd.isna(psnr_u) or pd.isna(psnr_v):
                        psnr_yuv611 = None
                    else:
                        psnr_yuv611 = (6 * psnr_y + psnr_u + psnr_v) / 8

                    all_data.append(
                        {
                            "Sequence": name,
                            "QP": qp,
                            "Rate": rate,
                            "PSNR-Y": psnr_y,
                            "PSNR-U": psnr_u,
                            "PSNR-V": psnr_v,
                            "PSNR-YUV": psnr_yuv611,
                        }
                    )
                except Exception as e:
                    print(f"Skipping row {i+1} in sheet '{sheet_name}': {e}")

            rdplot_df = pd.DataFrame(all_data)

            out_sheet_name = sheet_name.replace(".", "")
            out_name = f"{out_sheet_name}_rdplot.csv"
            rdplot_df.to_csv(out_name, index=False, sep=";")
            print(f"CSV exported: {out_name}")


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("excel_file", type=str, help="Input excel file")

    parser.add_argument(
        "--csv_format",
        type=str,
        choices=["normal", "rdplot"],
        help="format of output csv to use",
        default="normal",
    )

    pargs = parser.parse_args()

    file_path = pargs.excel_file
    csv_format = pargs.csv_format

    if not os.path.isfile(file_path):
        print(f"Error: File '{file_path}' not found")
        sys.exit(1)

    extract_all(file_path, csv_format)