# BrainGenix-NES
# AGPLv3

import json
import base64
import tqdm
import yaspin
import time

from . import Configuration

from BrainGenix.NES.Client import RequestHandler

import BrainGenix.LibUtils.ConfigCheck



class EM:

    def __init__(self, _Configuration:Configuration, _RequestHandler:RequestHandler, _SimulationID:int):
        # Create Attributes
        self.RequestHandler = _RequestHandler
        self.SimulationID = _SimulationID

        # Run Configuration Check
        BrainGenix.LibUtils.ConfigCheck.ConfigCheck(_Configuration)

        # Create On Server
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/Initialize?SimulationID={_SimulationID}")
        assert(Response != None)

        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/SetupMicroscope?SimulationID={_SimulationID}&PixelResolution_nm={_Configuration.PixelResolution_nm}&ImageWidth_px={_Configuration.ImageWidth_px}&ImageHeight_px={_Configuration.ImageHeight_px}&SliceThickness_nm={_Configuration.SliceThickness_nm}&ScanRegionOverlap_percent={_Configuration.ScanRegionOverlap_percent}&MicroscopeFOV_deg={_Configuration.MicroscopeFOV_deg}&NumPixelsPerVoxel_px={_Configuration.NumPixelsPerVoxel_px}")
        assert(Response != None)



    ## Access Methods
    def DefineScanRegion(self, _Point1_um:list, _Point2_um:list):
        Point1 = json.dumps(_Point1_um)
        Point2 = json.dumps(_Point2_um)
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/DefineScanRegion?SimulationID={self.SimulationID}&Point1_um={Point1}&Point2_um={Point2}")
        assert(Response != None)
        self.ScanRegionID = Response["ScanRegionID"]


    def QueueRenderOperation(self):
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/QueueRenderOperation?SimulationID={self.SimulationID}&ScanRegionID={self.ScanRegionID}")
        assert(Response != None)
        return Response["StatusCode"]


    def GetRenderStatus(self):
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/GetRenderStatus?SimulationID={self.SimulationID}")
        assert(Response != None)
        return Response
    

    def GetImageStack(self):
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/GetImageStack?SimulationID={self.SimulationID}&ScanRegionID={self.ScanRegionID}")
        assert(Response != None)
        return Response["RenderedImages"]


    def GetImage(self, _ImageHandle:str):
        Response = self.RequestHandler.MakeAuthenticatedQuery(f"/NES/VSDA/EM/GetImage?SimulationID={self.SimulationID}&ImageHandle={_ImageHandle}")
        assert(Response != None)
        return bytes(Response["ImageData"], 'utf-8')
    

    def WaitForRender(self):

        # Setup Status Information
        StatusInfo:dict = self.GetRenderStatus()


        # Perform Sanity Check On Render
        if (StatusInfo["RenderStatus"] < 3):
            print("Error during rendering, API reports:\n")
            print(StatusInfo)


        # Block With Queued Spinner
        QueueSpinner = yaspin.yaspin(text="Render Operation In Queue, Elapsed Time", color="green", timer=True)
        QueueSpinner.start()
        while (StatusInfo["RenderStatus"] == 3):

            # Get Status Info, Wait To Avoid API Spam
            StatusInfo:dict = self.GetRenderStatus()
            time.sleep(0.25)

        QueueSpinner.ok()


        
        # Setup Slice, Image Status Bar
        SliceStatusBar = tqdm.tqdm("Rendering Slice", total=1)
        SliceStatusBar.leave = True
        SliceStatusBar.colour = "green"
        LastSliceNumber = 0

        ImageStatusBar = tqdm.tqdm("Rendering Image", total=1)
        ImageStatusBar.leave = True
        ImageStatusBar.colour = "magenta"


        # Block Execution Until Render Finishes, Update Bar As We Wait
        while (StatusInfo["RenderStatus"] != 5):

            # Update Slice Bar
            SliceStatusBar.total = int(StatusInfo["TotalSlices"])
            SliceStatusBar.n = max(int(StatusInfo["CurrentSlice"]), 0)
            SliceStatusBar.refresh()
            SliceStatusBar.set_description(f"Rendering Slice {StatusInfo['CurrentSlice']}/{StatusInfo['TotalSlices']}")


            # Update Image Bar
            if (SliceStatusBar.n > LastSliceNumber):
                LastSliceNumber = SliceStatusBar.n
            ImageStatusBar.total = int(StatusInfo["TotalSliceImages"])
            ImageStatusBar.n = max(int(StatusInfo["CurrentSliceImage"]), 0)
            ImageStatusBar.refresh()
            ImageStatusBar.set_description(f"Rendering Image {StatusInfo['CurrentSliceImage']}/{StatusInfo['TotalSliceImages']}")



            # Get Status Info, Wait To Avoid API Spam
            StatusInfo:dict = self.GetRenderStatus()
            time.sleep(0.1)
        
    
        # When we're done, rendering is finished - make it look like it truly is, then close the bar
        SliceStatusBar.total = int(StatusInfo["TotalSlices"])
        SliceStatusBar.n = int(StatusInfo["TotalSlices"])
        SliceStatusBar.set_description(f"Rendering Slice {StatusInfo['TotalSlices']}/{StatusInfo['TotalSlices']}")
        SliceStatusBar.refresh()
        SliceStatusBar.close()
        ImageStatusBar.total = int(StatusInfo["TotalSliceImages"])
        ImageStatusBar.n = int(StatusInfo["TotalSliceImages"])
        ImageStatusBar.set_description(f"Rendering Image {StatusInfo['TotalSliceImages']}/{StatusInfo['TotalSliceImages']}")
        ImageStatusBar.refresh()
        ImageStatusBar.close()



    def SaveImageStack(self, _ImageStackDirectoryPrefix:str = ""):

        # Get Image Stack Manifest
        ImageHandles = self.GetImageStack()

        # Setup Progress Bar
        Bar = tqdm.tqdm("Downloading Image Stack", total=len(ImageHandles))
        Bar.leave = True
        Bar.colour = "green"


        for i in range(len(ImageHandles)):

            # Save This Image
            ImageData = self.GetImage(ImageHandles[i])
            with open(_ImageStackDirectoryPrefix + ImageHandles[i].split("/")[1], "wb") as FileHandler:
                FileHandler.write(base64.decodebytes(ImageData))

            Bar.set_description(f"Downloading Image {i + 1}/{len(ImageHandles)}")

            # Count Up Bar
            Bar.n = i + 1
            Bar.refresh()

        Bar.close()