from qolab.file_utils import save_table_with_header
class Trace:
    def __init__(self, label):
        self.label = label
        self.values = None
        self.unit = None
        self.tags = {}

    def plot(self):
        import matplotlib.pyplot as plt
        plt.plot(self.values, label=self.descr)
        plt.legend()
    
    def header(self):
        header = []
        if self.label is not None:
            header.append(f'label = {self.label}')
        if self.unit is not None:
            header.append(f'unit = {self.unit}')
        for k, v in self.tags.items():
            header.append(f'{k} = {v}')
        return header

    def write(self, fname):
        pass
        


class TraceXY:
    def __init__(self, label):
        self.label = label
        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)
        
class TraceSetSameX:
    def __init__(self, label):
        self.label = label
        self.x = None
        self.traces={}

    def addTrace(self, tr):
        if len(self.traces) == 0:
            self.x = tr.x
        trY = tr.y
        self.traces[tr.label]=trY
        
    def plot(self):
        import matplotlib.pyplot as plt
        for k, tr in self.traces.items():
            plt.plot(self.x.values, tr.values, label=tr.label)
        plt.xlabel(self.x.label)
        plt.legend()

    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)