import trsfile
from trsfile.parametermap import TraceSetParameterMap
import trsfile.traceparameter as tp
from trsfile.parametermap import TraceParameterMap, TraceParameterDefinitionMap
from trsfile.traceparameter import ByteArrayParameter, ParameterType, TraceParameterDefinition
import matplotlib.pyplot as plt
import sys
import numpy as np
#import random, os
import math 

parameters = TraceSetParameterMap()
print(parameters)

zoomS=0
zoomE=220000

start = 0
number = 1
displayLabels = 1

window = 100
overlap = 0.99
ABS=True
step = window - int(window * overlap)

with trsfile.open(sys.argv[1], 'r') as traces:
    print(traces.get_headers())
    # Show all headers
    for header, value in traces.get_headers().items():
        print(header, '=', value)
    scale_X = traces.get_headers().get(trsfile.Header.SCALE_X)
    scale_Y = traces.get_headers().get(trsfile.Header.SCALE_Y)
    lengthData = traces.get_headers().get(trsfile.Header.LENGTH_DATA)
    print(lengthData)

    if ABS:
        nameTrace = sys.argv[1][0:-4]+'+AWR('+str(window)+','+str(overlap)+').trs'
    else:
        nameTrace = sys.argv[1][0:-4]+'+WR('+str(window)+','+str(overlap)+').trs'
    with trsfile.trs_open(
        nameTrace,                 # File name of the trace set
        'w',                             # Mode: r, w, x, a (default to x)
        # Zero or more options can be passed (supported options depend on the storage engine)
        engine = 'TrsEngine',            # Optional: how the trace set is stored (defaults to TrsEngine)
        headers = {                      # Optional: headers (see Header class)
            trsfile.Header.TRS_VERSION: 2,
            trsfile.Header.SCALE_X: scale_X,
            trsfile.Header.SCALE_Y: scale_Y,
            trsfile.Header.DESCRIPTION: 'Window Resampled Traces',
            #trsfile.Header.TRACE_PARAMETER_DEFINITIONS: TraceParameterDefinitionMap(
            #    {'LEGACY_DATA': TraceParameterDefinition(ParameterType.BYTE, 16, 0)})
        },
        padding_mode = trsfile.TracePadding.AUTO,# Optional: padding mode (defaults to TracePadding.AUTO)
        live_update = True               # Optional: updates the TRS file for live preview (small performance hit)
                                         #   0 (False): Disabled (default)
                                         #   1 (True) : TRS file updated after every trace
                                         #   N        : TRS file is updated after N traces
    ) as wrtraces:

        for i, trace in enumerate(traces[(start):(start+number)]):
            print("here")
            trace_array = trace.samples[zoomS:zoomE]
            #print(len(trace_array))
            data = trace.parameters['LEGACY_DATA'].value
            print(len(data))
            length = math.ceil((zoomE - zoomS) / step)
            print(length)
            processed = np.zeros(length,np.float)

            for j in range(length):
                start = j*step
                if ABS:
                    chunk = np.abs(trace_array[start:(start+window)])
                else:
                    chunk = trace_array[start:(start+window)]
                processed[j]= np.sum(chunk)/len(chunk)
                #print(len(chunk))

            # Adding one Trace
            wrtraces.append(
                trsfile.Trace(
                    trsfile.SampleCoding.FLOAT,
                    #trace_array, 
                    processed,
                    TraceParameterMap({'LEGACY_DATA': ByteArrayParameter(data)})
                )
            )