import sys
import platform
import traceback
from BabelViscoFDTD.tools.RayleighAndBHTE import  InitOpenCL, InitCuda, InitMetal,InitMLX
from BabelViscoFDTD.H5pySimple import ReadFromH5py, SaveToH5py
from scipy.io import savemat
import numpy as np

from ThermalModeling.CalculateTemperatureEffects import CalculateTemperatureEffects
from multiprocessing import Process,Queue


class InOutputWrapper(object):
    def __init__(self, queue, stdout=True):
        self.queue=queue
        if stdout:
            self._stream = sys.stdout
            sys.stdout = self
        else:
            self._stream = sys.stderr
            sys.stderr = self
        self._stdout = stdout

    def write(self, text):
        self.queue.put(text)

    def __getattr__(self, name):
        return getattr(self._stream, name)

    def __del__(self):
        try:
            if self._stdout:
                sys.stdout = self._stream
            else:
                sys.stderr = self._stream
        except AttributeError:
            pass

#subprocess , useful to deal with large number of cases
def SubProcess(queueMsg,queueResult,case,deviceName,**kargs):
    stdout = InOutputWrapper(queueMsg,True)
    if kargs['Backend']=='CUDA':
        InitCuda(deviceName)
    elif kargs['Backend']=='OpenCL':
        InitOpenCL(deviceName)
    elif kargs['Backend']=='Metal':
        InitMetal(deviceName)
    elif kargs['Backend']=='MLX':
        InitMLX(deviceName)
    fname=CalculateTemperatureEffects(case,deviceName,queueMsg,**kargs)
    queueResult.put(fname)
    

def CalculateThermalProcess(queueMsg,case,AllDC_PRF_Duration,ExtraData,**kargs):

    try:
        Backend = ['CUDA','OpenCL','Metal','MLX'][kargs['COMPUTING_BACKEND']-1]
        deviceName=kargs['deviceName']
        AllCases=[]
        #These fields will preseved individually per sonication regime
        lf =['MaxBrainPressure','MaxIsppa', 'MaxIspta','TI','TIC','TIS','TempProfileTarget',\
            'TimeProfileTarget','p_map_central','Isppa','Ispta','MI','DurationUS','DurationOff',\
            'DutyCycle','PRF','BaselineTemperature','Repetitions','NumberGroupedSonications',\
            'PauseBetweenGroupedSonications']
        Index=[]
        bFirstCombination=True
        for combination in AllDC_PRF_Duration:
            SubData={}

            queueResult=Queue()
            kargsSub={}
            for k in ['PRF','DurationOff','Repetitions','NumberGroupedSonications','PauseBetweenGroupedSonications']:
                kargsSub[k]=combination[k]
            kargsSub['Isppa']=kargs['Isppa']
            kargsSub['sel_p']=kargs['sel_p']
            kargsSub['DutyCycle']=combination['DC']            
            kargsSub['DurationUS']=combination['Duration']
            kargsSub['bForceRecalc']=True
            kargsSub['Backend']=Backend
            kargsSub['Frequency']=kargs['Frequency']
            kargsSub['BaselineTemperature']=kargs['BaselineTemperature']
            kargsSub['LimitBHTEIterationsPerProcess']=kargs['LimitBHTEIterationsPerProcess']
            kargsSub['TxSystem']=kargs['TxSystem']
            if not bFirstCombination and kargs['bConcatenateSimulations']: #we concatenate sim results
                kargsSub['prevSimulationResultsFile']=prevSimulationResultsFile

            #we check now if some development parameters are being passed
            for k in ['bForceNoAbsorptionSkullScalp','bForceHomogenousMedium','HomogenousMediumValues','BenchmarkTestFile']:
                if k in kargs:
                    kargsSub[k]=kargs[k]
            fieldWorkerProcess = Process(target=SubProcess, 
                                    args=(queueMsg,queueResult,case,deviceName),
                                    kwargs=kargsSub)

            fieldWorkerProcess.start()
            fieldWorkerProcess.join()
            fname=queueResult.get()
            Data=ReadFromH5py(fname+'.h5')
            for f in lf:
                if 'p_map_central'==f:
                    SubData['p_map']=Data[f] #this will make it compatible for other purposes
                if f in Data: 
                    SubData[f]=Data[f]
            AllCases.append(SubData)
            Index.append([combination['DC'],combination['PRF'],combination['Duration'],combination['DurationOff'],np.round(kargs['Isppa'],1)])
            prevSimulationResultsFile=fname+'.h5'
            bFirstCombination=False
        for f in lf:
            if f in Data: #this will prevent failing with fields that are specific to phase arrays such as steering
                Data.pop(f)
        Index=np.array(Index)

        Data['AllData']=AllCases
        Data['Index']=Index
        for k in ExtraData:
            Data[k]=ExtraData[k]
        ConsolodidateName=fname.split('-Duration-')[0]+'_AllCombinations'
        savemat(ConsolodidateName+'.mat',Data)
        SaveToH5py(Data,ConsolodidateName+'.h5')

    except BaseException as e:
        print('--Babel-Brain-Low-Error')
        print(traceback.format_exc())
        print(str(e))