import json
from pathlib import Path
import numpy as np
mipnerf360_scenes = ["bicycle", "flowers", "garden", "stump", "treehill", "room", "counter", "kitchen", "bonsai"]
blender_scenes = ["chair", "drums", "ficus", "hotdog", "lego", "materials", "mic", "ship"]

dirs = [("orig", "../mip-splatting"), ("no3dopcomp", "../mip-splatting-no3dopacity")]

mipnerf360_benchmarks = [("benchmark_360v2_ours", "Multi-scale Training and Multi-scale Testing on the the Blender dataset"),
                   ("benchmark_360v2_ours_stmt", "Single-scale Training and Multi-scale Testing on the the Blender dataset")]
blender_benchmarks = [("benchmark_nerf_synthetic_ours_mtmt", "Multi-scale Training and Multi-scale Testing on the Mip-NeRF 360 dataset"),
                   ("benchmark_nerf_synthetic_ours_stmt", "Single-scale Training and Multi-scale Testing on the Mip-NeRF 360 dataset")  ]

metrics = [("PSNR", True), ("SSIM", True), ("LPIPS", False), ("Count", False)]

all_benchmarks = [(blender_benchmarks, blender_scenes), (mipnerf360_benchmarks, mipnerf360_scenes)]

for benchmarks, scenes in all_benchmarks:
    for benchmark, description in benchmarks:
        print(f"\n\n{benchmark}: {description}")
        for metric, higher_is_better in metrics:
            print(f"\n{metric}:")
            values = {}
            for header, dir in dirs:
                values[header] = {}
                sum = 0
                for scene in scenes:
                    if metric == "Count":
                        with open(Path(dir) / benchmark / scene / "point_cloud" / "iteration_30000" / "point_cloud.ply", "rb") as file:
                            # Skip the first two lines
                            for _ in range(2):
                                _ = file.readline()
                            
                            # Read the third line
                            line = file.readline().strip().split()
                            
                            assert line[0] == b"element" and line[1] == b"vertex"
                            val = int(line[2])
                    else:
                        with open(Path(dir) / benchmark / scene / "results.json", "r") as f:
                            data = json.load(f)
                            val = data["ours_30000"][metric]
                    values[header][scene] = val
                    sum += val
                values[header]["Average"] = sum / len(scenes)
            column_headers = scenes + ["Average"]
            rank = {}
            for header in values.keys():
                rank[header] = {}
            for scene in column_headers:
                column = []
                for header in values.keys():
                    column.append(values[header][scene])
                column = np.array(column)
                order = column.argsort()
                ranks = order.argsort()
                if higher_is_better:
                    ranks = len(ranks) - ranks - 1
                for i, header in enumerate(values.keys()):
                    rank[header][scene] = ranks[i]
            print("| |", end="")
            for scene in column_headers:
                print(f" {scene} |", end="")
            print("\n|-|"+len(column_headers)*"-|", end="")
            for header in values.keys():
                print(f"\n| {header} |", end="")
                for scene in column_headers:
                    highlight = ""
                    if rank[header][scene] == 0:
                        highlight = "**"
                    elif len(rank) > 2 and rank[header][scene] == 1:
                        highlight = "*"
                    print(f" {highlight}{values[header][scene]:.3f}{highlight} |", end="")
            print("")