cycles = ['LDCC', 'LDR', 'LRP']
variables = ['iBatt_10_s', 'uBatt_10_s', 'tCell_10_s', 'SOC_Cell_10_s']
save_var = ['i', 'u', 'T', 'SOC']

DATA = {}
for cycle in cycles:
    print('* Loading cycle in memory: ', cycle)
    for i, variable in enumerate(variables):
        print('    Loading variable in memory: ', variable)
        key_name =  cycle + '_' + variable
        target_csv = r'/Users/dreamer/Desktop/pytorchh/csv/MiL_sim_' + key_name + '.txt'
        f = open(target_csv, 'r')
        fresh_data = f.read()
        f.close()
        

        fresh_data = fresh_data.split(',')
        
        save_name = cycle + '_' + save_var[i]
        DATA[save_name] = np.array(fresh_data).astype(float)
        

# set their shapes right
for key in DATA:
    DATA[key] = DATA[key].reshape(DATA[key].shape[0], 1)


LDCC_i   = []
LDCC_u   = []
LDCC_T   = []
LDCC_SOC = []

k = 0
limit_of_jump = 20
wrt = DATA['LDCC_u']
for i in range(len(wrt) - 1):
    if wrt[i] == wrt[i + 1]:
        k = k + 1
    if k == limit_of_jump or wrt[i] != wrt[i + 1]:
        k = 0
        LDCC_SOC.append(DATA['LDCC_SOC'][i])
        LDCC_i.append(DATA['LDCC_i'][i])
        LDCC_u.append(DATA['LDCC_u'][i])
        LDCC_T.append(DATA['LDCC_T'][i])
LDR_i   = []
LDR_u   = []
LDR_T   = []
LDR_SOC = []
k = 0
limit_of_jump = 20
wrt = DATA['LDR_u']
for i in range(len(wrt) - 1):
    if wrt[i] == wrt[i + 1]:
        k = k + 1
    if k == limit_of_jump or wrt[i] != wrt[i + 1]:
        k = 0
        LDR_SOC.append(DATA['LDR_SOC'][i])
        LDR_i.append(DATA['LDR_i'][i])
        LDR_u.append(DATA['LDR_u'][i])
        LDR_T.append(DATA['LDR_T'][i])
def scaler(tensor):
    for ch in tensor:
        scale = 1.0 / (ch.max(dim=0)[0] - ch.min(dim=0)[0])        
        ch.mul_(scale).sub_(ch.min(dim=0)[0])        
    return tensor

filtered_LDCC_u = savgol_filter(np.ravel(LDCC_u), 501, 3) # 51, 3


train_i = np.concatenate((LDCC_i, LDR_i))
train_u = np.concatenate((LDCC_u, LDR_u))
train_T = np.concatenate((LDCC_T ,LDR_T))

train_inputs  = np.concatenate((train_i, train_u, train_T), axis=1)
train_outputs = np.concatenate((LDCC_SOC, LDR_SOC))

input_scaler = MinMaxScaler(feature_range=(-1,1))     
input_scaler.fit(train_inputs)

input_scaler_outputs=MinMaxScaler(feature_range=(-1,1))
input_scaler_outputs.fit(train_outputs)

inputs_LDCC  = input_scaler.transform(np.concatenate((LDCC_i, LDCC_u, LDCC_T), axis=1))
#input_scaler.fit(train_inputs)
#inputs_LDCC1  = input_scaler.fit_transform(np.concatenate((LDCC_i, LDCC_u, LDCC_T), axis=1))
outputs_LDCC=input_scaler_outputs.transform(np.array(LDCC_SOC))



#train_inputs=torch.Tensor(np.concatenate((LDCC_i, LDCC_u, LDCC_T), axis=1))
#inputs_LDCC1=scaler(train_inputs)




inputs_LDR  = input_scaler.transform(np.concatenate((DATA['LDR_i'], DATA['LDR_u'], DATA['LDR_T']), axis=1))
outputs_LDR = input_scaler_outputs.transform(DATA['LDR_SOC'])
print(inputs_LDR.shape)
inputs_LRP  = input_scaler.transform(np.concatenate((DATA['LRP_i'], DATA['LRP_u'], DATA['LRP_T']), axis=1))
outputs_LRP = input_scaler_outputs.transform(DATA['LRP_SOC'])

#inputs_LDCC1=inputs_LDCC1.numpy()

data_x = np.concatenate((inputs_LDCC, inputs_LDR), axis=0)

data_x = data_x.reshape(data_x.shape[0], 1, data_x.shape[1])
data_y = np.concatenate((outputs_LDCC, outputs_LDR), axis=0)
print("data y: ",data_y.shape)
val_range = int(data_x.shape[0]/100) * 15

val_x = data_x[0:val_range, :,:]
train_x = data_x[val_range:None, :,:]

val_y = data_y[0:val_range, :]
train_y = data_y[val_range:None, :]



'''
for axis in range(0, data_x.shape[1]):
    plt.figure()
    plt.plot(data_x[:, axis])
plt.show()

'''
train_x=torch.Tensor(train_x)
val_x=torch.Tensor(val_x)
print("train x size:",train_x.size())

class Model(nn.Module):
        def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv1d(1, 20, 1, 1)
                self.conv2 = nn.Conv1d(20, 50, 1, 1)
                self.fc1 = nn.Linear(1*3*50, 500)
                #self.dropout1 = nn.Dropout(0.5)
                self.fc2 = nn.Linear(500, 10)
        def forward(self, x):
                x = F.relu(self.conv1(x))
                x = F.max_pool1d(x, 1, 1)
                x = F.relu(self.conv2(x))
                x = F.max_pool1d(x, 1, 1)
                x = x.view(-1, 1*3*50)
                x = F.relu(self.fc1(x))
                #x = self.dropout1(x)
                x = self.fc2(x)
                return x
'''
class Model(nn.Module):
        def __init__(self,i,h1,out):
                super().__init__()
                self.linear=nn.Linear(i,h1)
                self.linear2=nn.Linear(h1,out)
        def forward(self,x):
                x=F.tanh(self.linear(x))
                x=F.tanh(self.linear2(x))
                return x


'''
'''
data_x = inputs_LDCC.reshape(inputs_LDCC.shape[0], 1, inputs_LDCC.shape[1])
data_y = outputs_LDCC

val_range = int(data_x.shape[0]/100) * 15

val_x   = data_x[0:val_range, :, :]
train_x = data_x[val_range:None, :, :]

val_y = data_y[0:val_range, :]
train_y = data_y[val_range:None, :]
train_x=torch.Tensor(train_x)
train_x=train_x.view(train_x.size(0),-1)
print(train_x.size(0))
val_x=torch.Tensor(val_x)
'''

batch_size=1000
epochs=15
device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu")
tr=torch.utils.data.DataLoader(train_x,batch_size=100,shuffle=True,drop_last=True)
val=torch.utils.data.DataLoader(val_x,batch_size=100,shuffle=False)
model=Model().to(device)
print(tr.__len__())
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

next(iter(tr))

for e in range(epochs):
  
        running_loss = 0.0
        running_corrects = 0.0
        val_running_loss = 0.0
        val_running_corrects = 0.0
  
        for inputs in tr:
                inputs = inputs.to(device)

                outputs = model(inputs)
                loss = criterion(outputs, inputs)
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()