import pandas as pd
import numpy as np


# In[11]:


d1 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_ADRS_CHILD_FU2-IMAGEN_DIGEST.csv")
d2 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_ANXDX_CHILD_FU2-IMAGEN_DIGEST.csv")
d3 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_AUDIT_CHILD_FU2-IMAGEN_DIGEST.csv")
d4 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_BIS_CHILD_FU2-IMAGEN_DIGEST.csv")
d5 = pd.read_csv("//zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_CAPE_CHILD_FU2-IMAGEN_DIGEST.csv")
d6 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_CSI_CHILD_FU2-IMAGEN_DIGEST.csv")
d7 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_CTQ_CHILD_FU2-IMAGEN_DIGEST.csv")
d8 = pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_EDEQ_CHILD_FU2-IMAGEN_DIGEST.csv")
d9= pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_JVQ_CHILD_FU2-IMAGEN_DIGEST.csv")
d10= pd.read_csv("/zi/home/sajad.rezaei/clip/clip/02_Full_Pipeline/data/psytools/IMAGEN-IMGN_TFEQ_CHILD_FU2-IMAGEN_DIGEST.csv")
data_df= pd.read_csv('/zi/home/sajad.rezaei/clip/clip/01_scripts/00_Text_processing/00_data/01_IMAGEN/00_data/IMAGEN_FU2.csv')


# In[12]:


d7["User code"] = d7["User code"].str.replace("-I", "-C")
df_list = [d1, d2, d3, d4, d5, d6, d7, d8, d9, d10]


# In[13]:


patterns = [
    "adrs", "ANXDX", "audit", "BIS", "CAPE42", "item",
    "CTQ", "EDEQ", "SCID", "IRI", "JVQ", "PAAQ",
    "RRS", "tci", "TFEQ", "User code"
]


# In[16]:


def process_dataframe(df, patterns, drop_zero_rows=False):
    df = df[df.columns[df.columns.str.contains('|'.join(patterns))]]
    df.loc[:,"User code"]= df.loc[:,"User code"].str.replace("-C", "")
    df.loc[:,"User code"]= df.loc[:,"User code"].str.lstrip("0")
   
    # Convert all elements to numeric, forcing non-numeric to NaN
    df_numeric = df.apply(pd.to_numeric, errors='coerce')

    # Replace NaN values with 0
    df_numeric.fillna(0, inplace=True)
   
    # Replace empty strings with 0
    df_numeric.replace("", 0, inplace=True)
   
    # remove the row with all zeros besiders the first column
    if drop_zero_rows:
        df_numeric = df_numeric.loc[(df_numeric.iloc[:, 1:] != 0).any(axis=1)]
    return df_numeric


# In[17]:


df_list_processed = [process_dataframe(df, patterns, drop_zero_rows=True) for df in df_list]


# In[19]:


#drop rows with NaN
df_list_processed = [df.dropna() for df in df_list_processed]


# In[20]:


#merge all df on User code
merged_df = df_list_processed[0]
for i, df in enumerate(df_list_processed, start=2):
    merged_df = pd.merge(merged_df, df, on='User code', suffixes=('', f'_df{i}'))


# In[21]:


#drop all row of data_df besicdes the fist and the last
data_df = data_df.drop(data_df.columns[1:-1], axis=1)
data_df.rename(columns={'Unnamed: 0': 'User code'}, inplace=True)


# In[22]:


#merge the merged_df with data_df
merged_df = pd.merge(data_df, merged_df, on='User code')


# In[23]:


#length before and after processing
print("Length: ", len(merged_df))


# In[61]:


import torch
from torchvision import datasets
from torchvision import transforms
import matplotlib.pyplot as plt
import pandas as pd
import load_helper as lh
from torch.utils.data import DataLoader, random_split
import importlib
from torchmetrics.classification import BinaryAccuracy, BinaryAUROC


df = merged_df.drop(columns=["User code", "FilePath"])
non_binary_columns = [col for col in df.columns if len(df[col].unique()) > 2]
df = pd.get_dummies(df, columns=non_binary_columns, drop_first=True)
autoencoder_dataset  = lh.AutoencoderDataset(dataframe=df)
total_size = len(autoencoder_dataset)
val_size = int(total_size * 0.25)  
train_size = total_size - val_size  

# Split the dataset
train_dataset, val_dataset = random_split(autoencoder_dataset, [train_size, val_size])

train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=512, shuffle=False)

#Metrics
train_accuracy = BinaryAccuracy()
valid_accuracy = BinaryAccuracy()
train_AUC = BinaryAUROC()
valid_AUC = BinaryAUROC()


# In[185]:


# Assuming 'loader' is your DataLoader
for batch_idx, (data, _) in enumerate(train_loader):
    print(f"Batch {batch_idx+1}:")
    print(f"Data shape: {data.shape}")
    # If you want to break after the first batch to just see one example
    break


# In[186]:


class AE(torch.nn.Module):
    def __init__(self, input_size):
        super().__init__()
       
        self.encoder = torch.nn.Sequential(
            torch.nn.Linear(input_size, 64),
            torch.nn.GELU(),
            torch.nn.BatchNorm1d(64),
            torch.nn.Linear(64, 64),
            torch.nn.GELU(),
            torch.nn.BatchNorm1d(64),
            torch.nn.Linear(64, 32),
            torch.nn.GELU(),
            torch.nn.Linear(32, 16),
            torch.nn.ReLU(),
        )
       
        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(16, 32),
            torch.nn.GELU(),
            torch.nn.BatchNorm1d(32),
            torch.nn.Linear(32, 64),
            torch.nn.GELU(),
            torch.nn.BatchNorm1d(64),
            torch.nn.Linear(64, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, input_size),
            torch.nn.Sigmoid()  
            )
    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

import torch.nn as nn

# class AE(torch.nn.Module):
#     def __init__(self, input_size):
#         super().__init__()

#         self.encoder = nn.Sequential(
#             nn.Conv1d(input_size, 64, kernel_size=3, padding=1),
#             nn.ReLU(),
#             nn.MaxPool1d(2),
#             nn.Conv1d(64, 32, kernel_size=3, padding=1),
#             nn.ReLU(),
#             nn.MaxPool1d(2),
#             nn.Flatten()
#         )

#         self.decoder = nn.Sequential(
#             nn.Linear(32 , 64),
#             nn.ReLU(),
#             nn.Linear(64, 128),
#             nn.ReLU(),
#             nn.Unflatten(1, (128, 1)),
#             nn.ConvTranspose1d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
#             nn.ReLU(),
#             nn.ConvTranspose1d(64, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
#             nn.Flatten(),

#             nn.Linear(input_size, input_size)  # Add a final linear layer to match input size
#         )

#     def forward(self, x):
#         x = x.unsqueeze(1)
#         encoded = self.encoder(x)
#         decoded = self.decoder(encoded)
#         return decoded.squeeze(1)





# In[187]:


# Model Initialization
model = AE(df.shape[1])

 
# Validation using MSE Loss function
loss_function = torch.nn.MSELoss()
 
# Using an Adam Optimizer with lr = 0.1
optimizer = torch.optim.Adam(model.parameters(),
                             lr = 1e-2,
                             weight_decay = 1e-3)




# In[188]:


epochs = 50
outputs = []
train_losses = []
val_losses = []

early_stopper = lh.EarlyStopper(patience=10, min_delta=0.01)
for epoch in range(epochs):
    model.train()
    for (data, _) in train_loader:
      data = data
      # Output of Autoencoder
      reconstructed = model(data)
      # Calculating the loss function
      loss = loss_function(reconstructed, data)
      # The gradients are set to zero,
      # the gradient is computed and stored.
      # .step() performs parameter update
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
       
      # Storing the losses in a list for plotting
      train_losses.append(loss.item())

    with torch.no_grad():
          train_auc = train_AUC(reconstructed, data)
        #   train_acc = train_accuracy(reconstructed, data)
        #   print(f"Train Accuracy on batch: {train_acc}")
        #   print(f"Train AUC on batch: {train_auc}")
    avg_train_loss = sum(train_losses) / len(train_losses)
    # Validation phase
    model.eval()  # Set model to evaluation mode
    with torch.no_grad():
        val_loss = 0
        count = 0
        for (data, _) in val_loader:  
            data = data
            reconstructed = model(data)
            val_loss = loss_function(reconstructed, data)
            val_losses.append(val_loss.item())
            valid_acc = valid_accuracy(reconstructed, data)
            valid_auc = valid_AUC(reconstructed, data)
            # print(f"Valid Accuracy on batch: {valid_acc}")
            # print(f"valid_auc on batch: {valid_auc}")
        avg_val_loss = sum(val_losses) / len(val_losses)

    print(f"Epoch {epoch+1}, Training Loss: {avg_train_loss:.4f}, Validation Loss: {avg_val_loss:.4f}")
    total_train_acc = train_accuracy.compute()
    total_train_auc = train_AUC.compute()
    total_val_acc = valid_accuracy.compute()
    total_val_auc = valid_AUC.compute()
    print(f"Accuracy on all data: Train {total_train_acc}, Valid: {total_val_acc}")
    print(f"AUC on all data: Train {total_train_auc}, Valid: {total_val_auc}")
    print("..................................................")
    train_accuracy.reset()
    train_AUC.reset()
    valid_accuracy.reset()
    valid_AUC.reset()

    if early_stopper.early_stop(avg_val_loss):
        print("Early stopping triggered.")
        break, give me the full code, change it to tensorflow