import os
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
import PytorchBoot.stereotype as stereotype

@stereotype.runner("strategy_generator")
class StrategyGenerator(Runner):
    def __init__(self, config):
        super().__init__(config)
        self.load_experiment("generate")
        
    def run(self):
        self.demo(seq=16,num=100)
    
    def create_experiment(self, backup_name=None):
        super().create_experiment(backup_name)
        output_dir = os.path.join(str(self.experiment_path), "output")
        os.makedirs(output_dir)
    
    def load_experiment(self, backup_name=None):
        super().load_experiment(backup_name)
    
    def demo(self, seq, num=100):
        import os
        from utils.data_load import DataLoadUtil
        from utils.reconstruction import ReconstructionUtil
        import numpy as np
        
        component = self.config["generate"][0]["component"] #r"C:\Document\Local Project\nbv_rec\output"
        data_dir = ConfigManager.get("datasets", "components", component, "root_dir")
        model_path = os.path.join(data_dir, f"sequence.{seq}\\world_points.txt")
        model_pts = np.loadtxt(model_path)
        
        output_dir = os.path.join(str(self.experiment_path), "output")
        pts_list = []
        for idx in range(0,num):
            path = DataLoadUtil.get_path(data_dir, seq, idx)
            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
            
            sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005)
            pts_list.append(sampled_point_cloud)
        
        sampled_model_pts = ReconstructionUtil.downsample_point_cloud(model_pts, 0.005)
        np.savetxt(os.path.join(output_dir,"sampled_model_points.txt"), sampled_model_pts)
        thre = 0.005
        
        useful_view, useless_view = ReconstructionUtil.compute_next_best_view_sequence(model_pts, pts_list, threshold=thre)
        print("useful:", useful_view)
        print("useless:", useless_view)
        
        selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, useful_view)
        downsampled_selected_full_views = ReconstructionUtil.downsample_point_cloud(selected_full_views, thre)
        np.savetxt(os.path.join(output_dir,"selected_full_views.txt"), downsampled_selected_full_views)
        
        limited_useful_view, limited_useless_view = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=thre, overlap_threshold=0.3)
        print("limited_useful:", limited_useful_view)
        print("limited_useless:", limited_useless_view)
        
        limited_selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, limited_useful_view)
        downsampled_limited_selected_full_views = ReconstructionUtil.downsample_point_cloud(limited_selected_full_views, thre)
        np.savetxt(os.path.join(output_dir,"selected_full_views_limited.txt"), downsampled_limited_selected_full_views)
        import json
        for idx, score in limited_useful_view:
            path = DataLoadUtil.get_path(data_dir, seq, idx)
            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
            print("saving useful view: ", idx, " | score: ", score)
            np.savetxt(os.path.join(output_dir,f"useful_view_{idx}.txt"), point_cloud)
        with open(os.path.join(output_dir,f"useful_view.json"), 'w') as f:
            json.dump(limited_useful_view, f)
        print("seq length: ", len(useful_view), "limited seq length: ", len(limited_useful_view))