"""
Created on Thu June 2 2022
@author: J.M. Algarín, MRILab, i3M, CSIC, Valencia
@email: josalggui@i3m.upv.es
@Summary: This class is able to sweep any parameter from any sequence
"""

import os
import sys
#*****************************************************************************
# Add path to the working directory
path = os.path.realpath(__file__)
ii = 0
for char in path:
    if (char=='\\' or char=='/') and path[ii+1:ii+14]=='PhysioMRI_GUI':
        sys.path.append(path[0:ii+1]+'PhysioMRI_GUI')
        sys.path.append(path[0:ii+1]+'marcos_client')
    ii += 1
#******************************************************************************
import numpy as np
import seq.mriBlankSeq as blankSeq
import pyqtgraph as pg              # To plot nice 3d images
from plotview.spectrumplot import SpectrumPlot
from plotview.spectrumplot import Spectrum3DPlot

class SweepImage(blankSeq.MRIBLANKSEQ):
    def __init__(self):
        super(SweepImage, self).__init__()
        # Input the parameters
        self.addParameter(key='seqName', string='SWEEPinfo', val='SWEEP')
        self.addParameter(key='seqNameSweep', string='Sequence', val='Noise', field='OTH')
        self.addParameter(key='parameter0', string='Parameter 0 X-axis', val='bw', field='OTH')
        self.addParameter(key='start0', string='Start point 0', val=30.0, field='OTH')
        self.addParameter(key='end0', string='End point 0', val=50.0, field='OTH')
        self.addParameter(key='nSteps0', string='Number of steps 0', val=5, field='OTH')
        self.addParameter(key='parameter1', string='Parameter 1 Y-axis', val='larmorFreq', field='OTH')
        self.addParameter(key='start1', string='Start point 1', val=3.0, field='OTH')
        self.addParameter(key='end1', string='End point 1', val=4.0, field='OTH')
        self.addParameter(key='nSteps1', string='Number of steps 1', val=5, field='OTH')

    def sequenceInfo(self):
        print(" ")
        print("Genera sweep sequence")
        print("Author: Dr. J.M. Algarín")
        print("Contact: josalggui@i3m.upv.es")
        print("mriLab @ i3M, CSIC, Spain")

    def sequenceTime(self):
        return(0)  # minutes, scanTime

    def sequenceRun(self, plotSeq=0):
        # Inputs
        seqName = self.mapVals['seqNameSweep']
        parameters = [self.mapVals['parameter0'], self.mapVals['parameter1']]
        start = [self.mapVals['start0'], self.mapVals['start1']]
        end = [self.mapVals['end0'], self.mapVals['end1']]
        nSteps = [self.mapVals['nSteps0'], self.mapVals['nSteps1']]

        # Sweep
        sampled = []
        parVector0 = np.linspace(start[0], end[0], nSteps[0]) # Create vector with parameters to sweep
        parVector1 = np.linspace(start[1], end[1], nSteps[1])
        seq = self.sequenceList[seqName] # Select the sequence that we want to sweep with modified parameters
        parMatrix = np.zeros((nSteps[0]*nSteps[1], 2))
        n = 0
        for step0 in range(nSteps[0]):
            for step1 in range(nSteps[1]):
                parMatrix[n, 0] = parVector0[step0]
                parMatrix[n, 1] = parVector1[step1]
                seq.mapVals[parameters[0]] = parVector0[step0]
                seq.mapVals[parameters[1]] = parVector1[step1]
                seq.sequenceRun(0)
                seq.sequenceAnalysis()
                if 'sampledCartesian' in seq.mapVals:
                    sampled.append(seq.mapVals['sampledCartesian']) # sampledCartesian is four column kx, ky, kz and S(kx, ky, kz)
                elif 'sampledPoint' in seq.mapVals:
                    sampled.append(seq.mapVals['sampledPoint'])
                else:
                    print('No signal to plot')
                    return 0
                n += 1

        self.seq = seq
        self.sampled = sampled
        return 0

    def sequenceAnalysis(self, obj=''):
        nPoints = np.array(self.seq.mapVals['nPoints'])
        nSteps = [self.mapVals['nSteps0'], self.mapVals['nSteps1']]
        start = [self.mapVals['start0'], self.mapVals['start1']]
        end = [self.mapVals['end0'], self.mapVals['end1']]
        parVector0 = np.linspace(start[0], end[0], nSteps[0])  # Create vector with parameters to sweep


        if 'sampledCartesian' in self.seq.mapVals:    # In case of images
            # Initialize data and image variables as zeros
            dataSteps = np.zeros((nSteps[0] * nSteps[1], nPoints[1], nPoints[0]), dtype=complex)
            imageSteps = dataSteps

            # Get axes in strings
            axes = self.seq.mapVals['axes']
            axesDict = {'x': 0, 'y': 1, 'z': 2}
            axesKeys = list(axesDict.keys())
            axesVals = list(axesDict.values())
            axesStr = ['', '', '']
            n = 0
            for val in axes:
                index = axesVals.index(val)
                axesStr[n] = axesKeys[index]
                n += 1

            # Generate k-space maps and images
            for step in range(nSteps[0]*nSteps[1]):
                data = self.sampled[step][:, 3]
                data = np.reshape(data, (nPoints[2], nPoints[1], nPoints[0]))
                image = np.fft.ifftshift(np.fft.ifftn(np.fft.ifftshift(data)))
                dataSteps[step, :, :] = data[int(nPoints[2]/2), :, :]
                imageSteps[step, :, :] = image[int(nPoints[2]/2), :, :]

            # Plot image
            image = np.abs(imageSteps)
            image = image / np.max(np.reshape(image, -1)) * 100
            image = Spectrum3DPlot(image,
                                   title="%s sweep images" % self.mapVals['seqNameSweep'],
                                   xLabel=axesStr[0],
                                   yLabel=axesStr[1])
            imageWidget = image.getImageWidget()

            # Plot k-space
            kSpace = np.log10(np.abs(dataSteps))
            kSpace = kSpace / np.max(np.reshape(kSpace, -1)) * 100
            kSpace = Spectrum3DPlot(kSpace,
                                   title="%s sweep k-spaces" % self.mapVals['seqNameSweep'],
                                   xLabel=axesStr[0],
                                   yLabel=axesStr[1])
            kSpaceWidget = image.getImageWidget()

            self.saveRawData()
            return([imageWidget, kSpaceWidget])

        elif 'sampledPoint' in self.seq.mapVals:  # In case of points (calibration sequences)
            image = np.zeros((nSteps[0], nSteps[1]), dtype=complex)
            n = 0
            for step0 in range(nSteps[0]):
                for step1 in range(nSteps[1]):
                    image[step0, step1] = self.sampled[n]
                    n += 1

            # Plot image
            if nSteps[1]>1: # If we sweep two parameters, show a map
                image = np.abs(image)
                map = image = image / np.max(np.reshape(image, -1)) * 100
                map = Spectrum3DPlot(np.abs(image),
                                     xLabel=self.seq.mapNmspc[self.mapVals['parameter0']],
                                     yLabel=self.seq.mapNmspc[self.mapVals['parameter1']],
                                     title='%s sweep' % self.mapVals['seqNameSweep'])
                mapWidget = map.getImageWidget()
            else: # If we sweep only one parameter, show a line plot
                image = np.reshape(image, -1)
                mapWidget = SpectrumPlot(xData=parVector0,
                                         yData=[np.abs(image)],
                                         legend=[''],
                                         xLabel=self.seq.mapNmspc[self.mapVals['parameter0']],
                                         yLabel='Output amplitude',
                                         title='%s sweep' % self.mapVals['seqNameSweep'])
                self.mapVals['sweepResult'] = [parVector0, np.abs(image)]
            self.saveRawData()
            return([mapWidget])

if __name__=='__main__':
    seq = SweepImage()
    seq.sequenceRun()
    seq.sequenceAnalysis()

