from multiprocessing import Process
import time
import datetime
import os
import json
import numpy as np
from celldetective.io import (
    extract_position_name,
    locate_segmentation_model,
    auto_load_number_of_frames,
    load_frames,
    _check_label_dims,
    _load_frames_to_segment,
)
from celldetective.utils import (
    _rescale_labels,
    _segment_image_with_stardist_model,
    _segment_image_with_cellpose_model,
    _prep_stardist_model,
    _prep_cellpose_model,
    _get_normalize_kwargs_from_config,
    extract_experiment_channels,
    _estimate_scale_factor,
    _extract_channel_indices_from_config,
    config_section_to_dict,
    _extract_nbr_channels_from_config,
    _get_img_num_per_channel,
)

from pathlib import Path, PurePath
from glob import glob
from shutil import rmtree
from tqdm import tqdm
import numpy as np
from csbdeep.io import save_tiff_imagej_compatible
from celldetective.segmentation import (
    segment_frame_from_thresholds,
    merge_instance_segmentation,
)
import gc
from art import tprint

import concurrent.futures


class BaseSegmentProcess(Process):

    def __init__(self, queue=None, process_args=None, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.queue = queue

        if process_args is not None:
            for key, value in process_args.items():
                setattr(self, key, value)

        tprint("Segment")

        # Experiment
        self.locate_experiment_config()

        print(f"Position: {extract_position_name(self.pos)}...")
        print("Configuration file: ", self.config)
        print(f"Population: {self.mode}...")
        self.instruction_file = os.sep.join(
            ["configs", f"segmentation_instructions_{self.mode}.json"]
        )

        self.read_instructions()
        self.extract_experiment_parameters()
        self.detect_movie_length()
        self.write_folders()

    def read_instructions(self):
        print("Looking for instruction file...")
        instr_path = PurePath(self.exp_dir, Path(f"{self.instruction_file}"))
        if os.path.exists(instr_path):
            with open(instr_path, "r") as f:
                _instructions = json.load(f)
                print(f"Measurement instruction file successfully loaded...")
                print(f"Instructions: {_instructions}...")
            self.flip = _instructions.get("flip", False)
        else:
            self.flip = False

    def write_folders(self):

        self.mode = self.mode.lower()
        self.label_folder = f"labels_{self.mode}"

        if os.path.exists(self.pos + self.label_folder):
            print("Erasing the previous labels folder...")
            rmtree(self.pos + self.label_folder)
        os.mkdir(self.pos + self.label_folder)
        print(f"Labels folder successfully generated...")

    def extract_experiment_parameters(self):

        self.spatial_calibration = float(
            config_section_to_dict(self.config, "MovieSettings")["pxtoum"]
        )
        self.len_movie = float(
            config_section_to_dict(self.config, "MovieSettings")["len_movie"]
        )
        self.movie_prefix = config_section_to_dict(self.config, "MovieSettings")[
            "movie_prefix"
        ]
        self.nbr_channels = _extract_nbr_channels_from_config(self.config)
        self.channel_names, self.channel_indices = extract_experiment_channels(
            self.exp_dir
        )

    def locate_experiment_config(self):

        parent1 = Path(self.pos).parent
        self.exp_dir = parent1.parent
        self.config = PurePath(self.exp_dir, Path("config.ini"))

        if not os.path.exists(self.config):
            print(
                "The configuration file for the experiment could not be located. Abort."
            )
            self.abort_process()

    def detect_movie_length(self):

        try:
            self.file = glob(self.pos + f"movie/{self.movie_prefix}*.tif")[0]
        except Exception as e:
            print(f"Error {e}.\nMovie could not be found. Check the prefix.")
            self.abort_process()

        len_movie_auto = auto_load_number_of_frames(self.file)
        if len_movie_auto is not None:
            self.len_movie = len_movie_auto

    def end_process(self):

        self.terminate()
        self.queue.put("finished")

    def abort_process(self):

        self.terminate()
        self.queue.put("error")


class SegmentCellDLProcess(BaseSegmentProcess):

    def __init__(self, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.check_gpu()

        # Model
        self.locate_model_path()
        self.extract_model_input_parameters()
        self.detect_channels()
        self.detect_rescaling()

        self.write_log()

        self.sum_done = 0
        self.t0 = time.time()

    def extract_model_input_parameters(self):

        self.required_channels = self.input_config["channels"]
        if "selected_channels" in self.input_config:
            self.required_channels = self.input_config["selected_channels"]

        self.target_cell_size = None
        if (
            "target_cell_size_um" in self.input_config
            and "cell_size_um" in self.input_config
        ):
            self.target_cell_size = self.input_config["target_cell_size_um"]
            self.cell_size = self.input_config["cell_size_um"]

        self.normalize_kwargs = _get_normalize_kwargs_from_config(self.input_config)

        self.model_type = self.input_config["model_type"]
        self.required_spatial_calibration = self.input_config["spatial_calibration"]
        print(
            f"Spatial calibration expected by the model: {self.required_spatial_calibration}..."
        )

        if self.model_type == "cellpose":
            self.diameter = self.input_config["diameter"]
            self.cellprob_threshold = self.input_config["cellprob_threshold"]
            self.flow_threshold = self.input_config["flow_threshold"]

    def write_log(self):

        log = f"segmentation model: {self.model_name}\n"
        with open(self.pos + f"log_{self.mode}.txt", "a") as f:
            f.write(f"{datetime.datetime.now()} SEGMENT \n")
            f.write(log)

    def detect_channels(self):

        self.channel_indices = _extract_channel_indices_from_config(
            self.config, self.required_channels
        )
        print(
            f"Required channels: {self.required_channels} located at channel indices {self.channel_indices}."
        )
        self.img_num_channels = _get_img_num_per_channel(
            self.channel_indices, int(self.len_movie), self.nbr_channels
        )

    def detect_rescaling(self):

        self.scale = _estimate_scale_factor(
            self.spatial_calibration, self.required_spatial_calibration
        )
        print(f"Scale: {self.scale}...")

        if self.target_cell_size is not None and self.scale is not None:
            self.scale *= self.cell_size / self.target_cell_size
        elif self.target_cell_size is not None:
            if self.target_cell_size != self.cell_size:
                self.scale = self.cell_size / self.target_cell_size

        print(f"Scale accounting for expected cell size: {self.scale}...")

    def locate_model_path(self):

        self.model_complete_path = locate_segmentation_model(self.model_name)
        if self.model_complete_path is None:
            print("Model could not be found. Abort.")
            self.abort_process()
        else:
            print(f"Model path: {self.model_complete_path}...")

        if not os.path.exists(self.model_complete_path + "config_input.json"):
            print(
                "The configuration for the inputs to the model could not be located. Abort."
            )
            self.abort_process()

        with open(self.model_complete_path + "config_input.json") as config_file:
            self.input_config = json.load(config_file)

    def check_gpu(self):

        if not self.use_gpu:
            os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

    def run(self):

        try:

            if self.model_type == "stardist":
                model, scale_model = _prep_stardist_model(
                    self.model_name,
                    Path(self.model_complete_path).parent,
                    use_gpu=self.use_gpu,
                    scale=self.scale,
                )

            elif self.model_type == "cellpose":
                model, scale_model = _prep_cellpose_model(
                    self.model_name,
                    self.model_complete_path,
                    use_gpu=self.use_gpu,
                    n_channels=len(self.required_channels),
                    scale=self.scale,
                )

            list_indices = range(self.len_movie)
            if self.flip:
                list_indices = reversed(list_indices)

            for i, t in enumerate(tqdm(list_indices, desc="frame")):

                f = _load_frames_to_segment(
                    self.file,
                    self.img_num_channels[:, t],
                    scale_model=scale_model,
                    normalize_kwargs=self.normalize_kwargs,
                )

                if self.model_type == "stardist":
                    Y_pred = _segment_image_with_stardist_model(
                        f, model=model, return_details=False
                    )

                elif self.model_type == "cellpose":
                    Y_pred = _segment_image_with_cellpose_model(
                        f,
                        model=model,
                        diameter=self.diameter,
                        cellprob_threshold=self.cellprob_threshold,
                        flow_threshold=self.flow_threshold,
                    )

                if self.scale is not None:
                    Y_pred = _rescale_labels(Y_pred, scale_model=scale_model)

                Y_pred = _check_label_dims(Y_pred, file=self.file)

                save_tiff_imagej_compatible(
                    self.pos
                    + os.sep.join([self.label_folder, f"{str(t).zfill(4)}.tif"]),
                    Y_pred,
                    axes="YX",
                )

                del f
                del Y_pred
                gc.collect()

                # Send signal for progress bar
                self.sum_done += 1 / self.len_movie * 100
                mean_exec_per_step = (time.time() - self.t0) / (i + 1)
                pred_time = (self.len_movie - (i + 1)) * mean_exec_per_step
                self.queue.put([self.sum_done, pred_time])

        except Exception as e:
            print(e)

        try:
            del model
        except:
            pass

        gc.collect()
        print("Done.")

        # Send end signal
        self.queue.put("finished")
        self.queue.close()


class SegmentCellThresholdProcess(BaseSegmentProcess):

    def __init__(self, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.equalize = False

        # Model

        self.load_threshold_config()
        self.extract_threshold_parameters()
        self.detect_channels()
        self.prepare_equalize()

        self.write_log()

        self.sum_done = 0
        self.t0 = time.time()

    def prepare_equalize(self):

        for i in range(len(self.instructions)):

            if self.equalize[i]:
                f_reference = load_frames(
                    self.img_num_channels[:, self.equalize_time[i]],
                    self.file,
                    scale=None,
                    normalize_input=False,
                )
                f_reference = f_reference[:, :, self.instructions[i]["target_channel"]]
            else:
                f_reference = None

            self.instructions[i].update({"equalize_reference": f_reference})

    def load_threshold_config(self):

        self.instructions = []
        for inst in self.threshold_instructions:
            if os.path.exists(inst):
                with open(inst, "r") as f:
                    self.instructions.append(json.load(f))
            else:
                print("The configuration path is not valid. Abort.")
                self.abort_process()

    def extract_threshold_parameters(self):

        self.required_channels = []
        self.equalize = []
        self.equalize_time = []

        for i in range(len(self.instructions)):
            ch = [self.instructions[i]["target_channel"]]
            self.required_channels.append(ch)

            if "equalize_reference" in self.instructions[i]:
                equalize, equalize_time = self.instructions[i]["equalize_reference"]
                self.equalize.append(equalize)
                self.equalize_time.append(equalize_time)

    def write_log(self):

        log = f"Threshold segmentation: {self.threshold_instructions}\n"
        with open(self.pos + f"log_{self.mode}.txt", "a") as f:
            f.write(f"{datetime.datetime.now()} SEGMENT \n")
            f.write(log)

    def detect_channels(self):

        for i in range(len(self.instructions)):

            self.channel_indices = _extract_channel_indices_from_config(
                self.config, self.required_channels[i]
            )
            print(
                f"Required channels: {self.required_channels[i]} located at channel indices {self.channel_indices}."
            )
            self.instructions[i].update({"target_channel": self.channel_indices[0]})
            self.instructions[i].update({"channel_names": self.channel_names})

        self.img_num_channels = _get_img_num_per_channel(
            np.arange(self.nbr_channels), self.len_movie, self.nbr_channels
        )

    def parallel_job(self, indices):

        try:

            for t in tqdm(
                indices, desc="frame"
            ):  # for t in tqdm(range(self.len_movie),desc="frame"):

                # Load channels at time t
                masks = []
                for i in range(len(self.instructions)):
                    f = load_frames(
                        self.img_num_channels[:, t],
                        self.file,
                        scale=None,
                        normalize_input=False,
                    )
                    mask = segment_frame_from_thresholds(f, **self.instructions[i])
                    # print(f'Frame {t}; segment with {self.instructions[i]=}...')
                    masks.append(mask)

                if len(self.instructions) > 1:
                    mask = merge_instance_segmentation(masks, mode="OR")

                save_tiff_imagej_compatible(
                    os.sep.join(
                        [self.pos, self.label_folder, f"{str(t).zfill(4)}.tif"]
                    ),
                    mask.astype(np.uint16),
                    axes="YX",
                )

                del f
                del mask
                gc.collect()

                # Send signal for progress bar
                self.sum_done += 1 / self.len_movie * 100
                mean_exec_per_step = (time.time() - self.t0) / (
                    self.sum_done * self.len_movie / 100 + 1
                )
                pred_time = (
                    self.len_movie - (self.sum_done * self.len_movie / 100 + 1)
                ) * mean_exec_per_step
                self.queue.put([self.sum_done, pred_time])

        except Exception as e:
            print(e)

        return

    def run(self):

        self.indices = list(range(self.img_num_channels.shape[1]))
        if self.flip:
            self.indices = np.array(list(reversed(self.indices)))

        chunks = np.array_split(self.indices, self.n_threads)

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=self.n_threads
        ) as executor:
            results = results = executor.map(
                self.parallel_job, chunks
            )  # list(map(lambda x: executor.submit(self.parallel_job, x), chunks))
            try:
                for i, return_value in enumerate(results):
                    print(f"Thread {i} output check: ", return_value)
            except Exception as e:
                print("Exception: ", e)

        print("Done.")
        # Send end signal
        self.queue.put("finished")
        self.queue.close()
