import multiprocessing
from multiprocessing.sharedctypes import Value as mpValue
import ctypes
import queue
import time
import psutil

from pyrateshield import labels, Logger

from pyrateshield.radtracer import radtracer
from pyrateshield.pyshield.engine import Engine
import pandas as pd
LOG_LEVEL = Logger.LEVEL_DEBUG


CRITICAL_POINT_NAME         = 'Critical Point Name'
RADTRACER_DOSE              = 'RadTracer Dose [mSv]'
PYSHIELD_DOSE               = 'PyShield Dose [mSv]'
OCCUPANCY_FACTOR            = 'Occupancy factor'
RADTRACER_DOSE_CORRECTED    = 'RadTracer Dose corrected for occupancy [mSv]'
PYSHIELD_DOSE_CORRECTED     = 'PyShield Dose corrected for occupancy [mSv]'
SOURCE_NAME                 = 'Source Name'


class Worker(multiprocessing.Process, Logger):
    def __init__(self, project_queue, source_queue, dosemap_queue, update_flag, stop_flag):
        multiprocessing.Process.__init__(self)
        Logger.__init__(self, log_level=LOG_LEVEL)

        self.project_queue = project_queue
        self.source_queue = source_queue
        self.dosemap_queue = dosemap_queue
        self.update_flag = update_flag
        self.stop_flag = stop_flag        
    
    def update_project(self):        
        if self.project.dosemap.engine == labels.PYSHIELD:
            #self.barriers = pyshield.Barriers.from_pyrateshield(self.project)
            self.pyshield_engine = Engine.from_pyrateshield(self.project)
    
    def get_dosemap(self, source):
        if self.project.dosemap.engine == labels.PYSHIELD:
            return self.pyshield_engine.source_dosemap(source)
            # return pyshield.Sources.from_pyrateshield(
            #             model=self.project,
            #             source=source, 
            #             barriers=self.barriers,
            #             ).get_dose_map()
        elif self.project.dosemap.engine == labels.RADTRACER:
            return radtracer.dosemap_single_source(source, self.project)
            
    def run(self):
        while not self.stop_flag.value:            
            try:
                sources = self.source_queue.get(timeout=0.1)
            except queue.Empty:
                sources = None
            
            if self.update_flag.value:
                self.project = self.project_queue.get()
                self.update_project()
                self.update_flag.value = False
                self.logger.info(f'{self.name}, Updated project')
        
            if sources is not None:                
                for source in sources:                    
                    self.dosemap_queue.put( self.get_dosemap(source) )
                    
        self.logger.info(f'{self.name}, Stopped')


class Dosemapper(Logger):
    def __init__(self, processes=None):
        Logger.__init__(self, log_level=LOG_LEVEL)
        self.source_queue = multiprocessing.Queue()
        self.dosemap_queue = multiprocessing.Queue()
        self._workers = []
        
        self.max_cpu_count = psutil.cpu_count(logical=True)
        if processes is None:
            processes = psutil.cpu_count(logical=False)
        self.set_workers(processes)
    
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self._stop_workers()
    
    def update_project(self, project):
        # Make sure all worker processes are done updating from a previous call
        # to this method before updating again.
        while any( w.update_flag.value for w in self._workers ):
            time.sleep(0.1)
        
        for w in self._workers:
            w.project_queue.put(project)
            w.update_flag.value = True
        
    def get_dosemap(self, project, sources=None):
        # Process all sources by default
        if sources is None:
            sources = list(project.sources_nm) + list(project.sources_ct) + list(project.sources_xray)
        
        # Exclude sources which are not set to Enabled
        sources = [src for src in sources if src.enabled]
        
        # Check if there is any work to be done
        if not len(sources):
            return None
        
                
        #This is a bit of a hack to ensure PyShield only receives NM sources,
        #and that sources with the same position are processed by the same worker.
        #This allows PyShield to benefit from caching.
        sources_non_nm = [src for src in sources if src.label != labels.SOURCES_NM]
        if project.dosemap.engine == labels.PYSHIELD:
            src_pos_dct = {}
            sources_nm = [src for src in sources if src not in sources_non_nm]
            for source in sources_nm:
                # Make a dictionary of source_positions and group the corresponding sources
                # I.e.: { (x,y,z): [src1, src2, src3 ], (x,y,z): [src4, src5, ... ] }
                src_pos_dct.setdefault(tuple(source.position), []).append(source)
            sources_grouped = [srcs for srcs in src_pos_dct.values()]
        else:
            sources_grouped = [[src] for src in sources]
        
                       
        self.update_project(project)
        
        for grp in sources_grouped:
            self.source_queue.put(grp)
        
        # Collect the results. Make sure to get as many results from the queue
        # as the number of sources
        nr_of_sources = sum(len(grp) for grp in sources_grouped)
                
        dosemap = None        
        update_dosemap = lambda dm: dm if dosemap is None else (dosemap + dm)
        
        for i in range(nr_of_sources):
            dosemap = update_dosemap( self.dosemap_queue.get() )        
        
        if project.dosemap.engine == labels.PYSHIELD and len(sources_non_nm):
            # Temporarily switch to RadTracer to process the non-NM sources
            print(f"Switching to RadTracer for {len(sources_non_nm)} non-NM sources")
            project.dosemap.engine = labels.RADTRACER
            dosemap = update_dosemap( self.get_dosemap(project, sources=sources_non_nm) )
            print("Done... Switching back to pyShield")
            project.dosemap.engine = labels.PYSHIELD
        
        return dosemap        
    
    
    @staticmethod
    def empty_critical_point_result():
        return pd.DataFrame([{CRITICAL_POINT_NAME:      None,
                              SOURCE_NAME:              None,
                              RADTRACER_DOSE:           None,
                              PYSHIELD_DOSE:            None,
                              OCCUPANCY_FACTOR:         None,
                              RADTRACER_DOSE_CORRECTED: None,
                              PYSHIELD_DOSE_CORRECTED:  None
                              }])
    @staticmethod
    def get_source_critical_point_result(source, 
                                         critical_point, 
                                         dose_radtracer, 
                                         dose_pyshield):
        
        factor = critical_point.occupancy_factor
        
        return {CRITICAL_POINT_NAME:        critical_point.name,
                SOURCE_NAME:                source.name,
                RADTRACER_DOSE:             dose_radtracer,
                PYSHIELD_DOSE:              dose_pyshield,
                OCCUPANCY_FACTOR:           factor,
                RADTRACER_DOSE_CORRECTED:   factor * dose_radtracer,
                PYSHIELD_DOSE_CORRECTED:    factor * dose_pyshield}
                              
        

    def get_sources_critical_points(self, project):
        sources = list(project.sources_nm) + list(project.sources_ct) + list(project.sources_xray)
        
        # Exclude sources which are not set to Enabled
        sources = [src for src in sources if src.enabled]
        
        # Exclude points which are not set to Enabled
        crit_points = [crp for crp in project.critical_points if crp.enabled]
                
        if len(crit_points) == 0:            
            return self.empty_critical_point_result()
        
        reports = []        
        
        pyshield_engine = Engine.from_pyrateshield(project)
        
        for crit_point in crit_points:
            for source in sources:
                dose_radtracer = radtracer.pointdose_single_source(
                    crit_point.position, source, project)
                
                if source.label == labels.SOURCES_NM:
                    dose_pyshield = pyshield_engine.dose_at_point(crit_point.position,
                                                         sources=[source])
                else:
                    dose_pyshield = dose_radtracer

                reports += [self.get_source_critical_point_result(
                    source, crit_point, dose_radtracer, dose_pyshield)]
                                                      
        return pd.DataFrame(reports)
        
    
    def get_critical_points(self, project):
        report = self.get_sources_critical_points(project)
        summed_report = pd.pivot_table(report, index=CRITICAL_POINT_NAME, aggfunc=sum)
        
        # preserve order
        pnames = [crit_point.name for crit_point in project.critical_points]
        summed_report = summed_report.reindex(pnames)
        
        summed_report = summed_report.reset_index()
        
        return summed_report

    @property
    def n_workers(self):
        return len(self._workers)
    
    
    def set_workers(self, processes):
        if processes > self.max_cpu_count:
            self.logger.info(f"Max nr of parallel processes: {self.max_cpu_count}")
            processes = self.max_cpu_count
        if processes < 0:
            processes = 0
        
        
        if processes > self.n_workers:
            self._start_workers(processes - self.n_workers)
        elif processes < self.n_workers:
            self._stop_workers(self.n_workers - processes)
    
    
    def _start_workers(self, processes=None):
        self.logger.info(f"Starting {processes} workers")
        
        for i in range(processes):
            project_queue = multiprocessing.Queue()
            stop_flag = mpValue(ctypes.c_bool, False)
            update_flag = mpValue(ctypes.c_bool, False)
            args = (
                project_queue,
                self.source_queue,
                self.dosemap_queue,
                update_flag,
                stop_flag,
            )
            p = Worker(*args)            
            self._workers.append(p)
            p.start()
        
        
    def _stop_workers(self, processes=None):
        if processes is None:
            processes = self.n_workers
        
        stopworkers = [self._workers.pop() for i in range(processes)]
                    
        for w in stopworkers:
            w.stop_flag.value = True
        for w in stopworkers: 
            w.join()
        time.sleep(0.1)
        for w in stopworkers:
            if w.is_alive():
                self.logger.info("Stil alive")
            w.close()
        
            
if __name__ == "__main__":
    from pyrateshield.model import Model
    import timeit
    import time
    #model = Model.load_from_project_file('/Users/marcel/git/pyrateshield/example_projects/LargeProject/large_project.psp')
    #model = Model.load_from_project_file('../example_projects/LargeProject/large_project.psp')
    model = Model.load_from_project_file('../example_projects/SmallProject/project.psp')
    model = Model.load_from_project_file('../example_projects/Lu-177.psp')
    model.dosemap.grid_matrix_size = 120
    
    with Dosemapper() as dm:        
        model.dosemap.engine = labels.PYSHIELD
        print("PyShield", timeit.timeit(lambda: dm.get_dosemap(model), number=1) )
    
        time.sleep(1) 
        print()
        
        model.dosemap.engine = labels.RADTRACER
        print("Radtracer", timeit.timeit(lambda: dm.get_dosemap(model), number=1) )

    # from pyshield import Sources
    # sources = Sources.from_pyrateshield(model)
    # print("Pyshield single CPU:", end=" ", flush=True)
    # print(timeit.timeit(lambda: sources.get_dosemap(), number=1) )
    
    
    
    

