import mpi4py.MPI as MPI
import os
import yaml
from glob import glob

from injection_pipeline.SingleEpochImage import SingleEpochImage
from injection_pipeline.InjectionCatalog import InjectionCatalog
from injection_pipeline.Catalog.C6_SimCat import SimCat
from injection_pipeline.config import parse_args


class InjectionPipeline(object):

    def __init__(self, config_file=None):
        # Load configuration
        if config_file is None:
            args = parse_args()
            config_file = args.config_file

        with open(config_file, "r") as stream:
            try:
                self.config = yaml.safe_load(stream)
                for key, value in self.config.items():
                    print(key + " : " + str(value))
            except yaml.YAMLError as exc:
                print(exc)

        self.input_dir = self.config['input_dir']
        self.output_dir = self.config['output_dir']
        self.pointing_label_list = self.config['pointing_label_list']
        self.chip_label_list = self.config['chip_label_list']

    def genearte_path_list_for_one_pointing(self,
                                            input_dir,
                                            pointing_label,
                                            chip_label_list=None):
        """_summary_

        Args:
            input_dir (_type_): _description_
            pointing_label (_type_): _description_
            chip_label_list (_type_, optional): _description_. Defaults to None.

        Returns:
            _type_: _description_
        """
        pointing_dir = os.path.join(input_dir, pointing_label)
        if chip_label_list is None:
            image_path_list = glob(
                pointing_dir + '/CSST_MSC_MS_SCIE_*_' + '*_img_*')
        else:
            image_path_list = []
            for chip_label in chip_label_list:
                image_path = glob(pointing_dir + '/CSST_MSC_MS_SCIE_*_' +
                                  chip_label + '_img_*')[0]
                image_path_list.append(image_path)

        return image_path_list

    def run_pointing_list(self,
                          Catalog,
                          input_dir=None,
                          pointing_label_list=None,
                          output_dir=None,
                          chip_label_list=None):
        if input_dir is None:
            input_dir = self.input_dir
        if pointing_label_list is None:
            pointing_label_list = self.pointing_label_list
        if output_dir is None:
            output_dir = self.output_dir
        if chip_label_list is None:
            chip_label_list = self.chip_label_list
        
        image_path_list = []
        output_path_list = []
        try:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
        except OSError:
            pass

        for pointing_label in pointing_label_list:
            output_pointing_dir = os.path.join(output_dir, pointing_label)
            try:
                if not os.path.exists(output_pointing_dir):
                    os.makedirs(output_pointing_dir)
            except OSError:
                pass
            temp_img_path_list = self.genearte_path_list_for_one_pointing(
                input_dir=input_dir,
                pointing_label=pointing_label,
                chip_label_list=chip_label_list)
            image_path_list = image_path_list + temp_img_path_list
            output_path_list = output_path_list + \
                [output_pointing_dir] * len(temp_img_path_list)

        comm = MPI.COMM_WORLD
        ind_thread = comm.Get_rank()
        num_thread = comm.Get_size()

        for i in range(len(image_path_list)):
            if i % num_thread != ind_thread:
                continue
            image_path = image_path_list[i]
            output_path = output_path_list[i]
            image = SingleEpochImage(
                config=self.config, input_img_path=image_path, output_dir=output_path)
            inject_cat = InjectionCatalog(image=image)
            inject_cat.generate_positions(config=self.config)
            print("number of galaxies to be injected: %d" %
                  (inject_cat.nobjects))
            input_cat = Catalog(config=self.config,
                                chip=image.chip,
                                filt=image.filt,
                                nobjects=inject_cat.nobjects,
                                logger=image.chip_output.logger)
            image.inject_objects(pos=inject_cat.pos, cat=input_cat)
            image.save_injected_img()

def main():
    pipeline = InjectionPipeline()
    pipeline.run_pointing_list(Catalog=SimCat)


if __name__ == "__main__":
    # input_dir = "/public/home/fangyuedong/project/50sqDeg_L1_outputs/50sqDeg_Photo_W2/"
    # # pointing_label_list = ["MSC_0000000", "MSC_0000001",
    # #                        "MSC_0000002", "MSC_0000003", "MSC_0000004", "MSC_0000005", "MSC_0000006", "MSC_0000007", "MSC_0000008", "MSC_0000009"]
    # chip_label_list = None
    # pointing_label_list = ["MSC_0000000"]
    # chip_label_list = ["07"]
    # # output_dir = "/public/home/fangyuedong/project/injected_50sqDeg_L1_outputs/50sqDeg_Photo_W3/"
    # output_dir = "/public/home/fangyuedong/project/test_photometry/50sqDeg_Photo_W2/"
    # config_file = "/public/home/fangyuedong/project/injection_pipeline/config/config_injection.yaml"

    # pipeline = InjectionPipeline()
    # pipeline.run_pointing_list(Catalog=SimCat)
    main()
