import argparse, json
from pathlib import Path
import torch
from arguments import ModelParams, PipelineParams, OptimizationParams
from scene import Scene, FourierModel

parser = argparse.ArgumentParser()
lp = ModelParams(parser); op = OptimizationParams(parser); pp = PipelineParams(parser)
parser.add_argument('--load_iteration', type=int, default=6000)
args = parser.parse_args()
dataset = lp.extract(args)
gaussians = FourierModel(dataset.sh_degree, fourier_K=dataset.fourier_K,
                         use_beta_color=dataset.beta_color, sb_lobes=dataset.sb_lobes)
scene = Scene(dataset, gaussians, load_iteration=args.load_iteration, shuffle=False)
out = Path(dataset.model_path) / 'lod_models'
out.mkdir(parents=True, exist_ok=True)
full_r = gaussians._radii_logits.detach().clone()
full_p = gaussians._phases_raw.detach().clone()
full_k = gaussians.fourier_K
rows = []
for k in range(1, full_k+1):
    gaussians._radii_logits = torch.nn.Parameter(full_r[:, :k].clone().requires_grad_(False))
    gaussians._phases_raw = torch.nn.Parameter(full_p[:, :k].clone().requires_grad_(False))
    gaussians.fourier_K = k
    gaussians.active_K = k
    path = out / f'fourier_k{k}.ply'
    gaussians.save_ply(str(path))
    rows.append({'K':k,'path':str(path),'bytes':path.stat().st_size,'primitives':int(gaussians.get_xyz.shape[0])})
gaussians._radii_logits = torch.nn.Parameter(full_r)
gaussians._phases_raw = torch.nn.Parameter(full_p)
gaussians.fourier_K = full_k
gaussians.active_K = full_k
(out/'sizes.json').write_text(json.dumps(rows,indent=2))
print(json.dumps(rows,indent=2))