from qolab.file_utils import save_table_with_header
import numpy as np
import yaml

class Trace:
    def __init__(self, label):
        self.config = {}
        self.config['label'] = label
        self.config['unit'] = None
        self.config['tags'] = {}
        self.values = None

    def plot(self):
        import matplotlib.pyplot as plt
        plt.plot(self.values, label=self.descr)
        plt.legend()
        plt.grid()
    
    def getConfig(self):
        return( self.config )

    def getHeader(self, prefix=''):
        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
        header = header.split('\n')
        prefixed_header = [prefix+l for l in header]
        return prefixed_header

    def save(self, fname, item_format='e', header=[], **kwargs):
        header.extend(self.getHeader())
        save_table_with_header(fname, self.values, header, item_format=item_format, **kwargs)


class TraceXY:
    def __init__(self, label):
        self.config = {}
        self.config['label'] = label
        self.config['tags'] = {}
        self.x = None
        self.y = None

    def plot(self):
        import matplotlib.pyplot as plt
        plt.plot(self.x.values, self.y.values, label=self.label)
        plt.legend()
        plt.xlabel(self.x.label)
        plt.grid()

    def getConfig(self):
        config = self.config.copy()
        config['TraceX'] = {}
        config['TraceX'] = self.x.getConfig()
        config['TraceY'] = {}
        config['TraceY'] = self.y.getConfig()
        return( config )

    def getHeader(self, prefix=''):
        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
        header = header.split('\n')
        prefixed_header = [prefix+l for l in header]
        return prefixed_header
        
    def save(self, fname, item_format='e', header=[], **kwargs):
        data=np.concatenate((self.x.values, self.y.values), 1)
        header.extend(self.getHeader())
        save_table_with_header(fname, data, header, item_format=item_format, **kwargs)

class TraceSetSameX:
    def __init__(self, label):
        self.config = {}
        self.config['label'] = label
        self.config['tags'] = {}
        self.x = None
        self.traces={}

    def addTrace(self, tr):
        if len(self.traces) == 0:
            self.x = tr.x
        trY = tr.y
        self.traces[tr.config['label']]=trY
        
    def plot(self):
        import matplotlib.pyplot as plt
        nplots = len(self.traces.keys())
        cnt=0
        for k, tr in self.traces.items():
            cnt+=1
            if cnt == 1:
                ax1=plt.subplot(nplots, 1, cnt)
            else:
                plt.subplot(nplots, 1, cnt, sharex=ax1)
            plt.plot(self.x.values, tr.values, label=tr.label)
            plt.xlabel(self.x.label)
            plt.legend()
            plt.grid()

    def items(self):
        return (self.traces.items())

    def keys(self):
        return (self.traces.keys())

    def getTrace(self, label):
        tr = TraceXY(label)
        tr.x = self.x
        tr.y = self.traces[label]
        return (tr)
    
    def getConfig(self):
        config = self.config.copy()
        config['TraceX'] = {}
        config['TraceX'] = self.x.getConfig()
        config['TraceY'] = {}
        for k,v in self.traces.items():
            config['TraceY'][k] = v.getConfig()
        return( config )

    def getHeader(self, prefix=''):
        header = yaml.dump(self.getConfig(), default_flow_style=False, sort_keys=False)
        header = header.split('\n')
        prefixed_header = [prefix+l for l in header]
        return prefixed_header
        
    def save(self, fname, item_format='e', header=[], **kwargs):
        data=self.x.values
        for k,v in self.traces.items():
            data=np.concatenate((data, v.values), 1)
        header.extend(self.getHeader())
        save_table_with_header(fname, data, header, item_format=item_format, **kwargs)


if __name__ == '__main__': 
    print("Testing trace")
    x=Trace('x trace')
    x.values = np.random.normal(2,2,(4,1))
    x.values = np.array(x.values, int)
    x.config['unit']='s'
    x.config['tags']['tag1'] = 'xxxx'
    x.config['tags']['tag2'] = 'xxxx'
    x.save('xtrace.dat', skip_headers_if_file_exist=True)
    # print(x.getHeader())
    y=Trace('y trace')
    y.values = np.random.normal(2,2,(4,1))
    y.config['unit']='V'
    y.config['tags']['ytag2'] = 'yyyy'
    xy=TraceXY('xy trace')
    xy.config['tags']['xy tag']= 'I am xy tag'
    xy.x = x
    xy.y = y
    xy.save('xytrace.dat')
    # print(xy.getHeader())
    xyn = TraceSetSameX('many ys trace')
    xyn.config['tags']['descr'] = 'I am many ys trace'
    xy.config['label']='y1'
    xyn.addTrace(xy)
    xy.config['label']='y2'
    xyn.addTrace(xy)
    xy.config['label']='y3'
    xyn.addTrace(xy)
    xyn.save('xyntrace.dat')
    # print(xyn.getHeader())