# -- Code Cell --
import pandas as pd

# -- Code Cell --
comp = pd.read_csv("./train_data/composers.csv")
train = pd.read_csv('./train_data/train_data.csv')

# -- Code Cell --
comp.head()

# -- Code Cell --
train['id'].unique()

# -- Code Cell --
comp['datapointID'].unique()

# -- Code Cell --
df = train.groupby("id").agg(list) #aici am toate melodile

# -- Code Cell --
df.info()

# -- Code Cell --
comp_usefull = comp.groupby('datapointID').agg(list)
comp_usefull = comp_usefull.drop(columns='subtaskID')

# -- Code Cell --
df["answer"] = comp_usefull["answer"].apply(lambda x: x[0])

# -- Code Cell --
df

# -- Code Cell --
df['answer'].unique()

# -- Code Cell --
mapping = {
    'Led Zeppelin': 0,
    'David Bowie': 1,
    'Psychedelic Corn Cruumpets': 2,
    'The Beatles': 3,
    'Pantera': 4
}

df['answer'] = df['answer'].map(mapping)

# -- Code Cell --
df

# -- Code Cell --
from catboost import CatBoostClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import f1_score
X = df.drop(columns='answer')
y = df['answer']
X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.2, random_state=42)
model = CatBoostClassifier(random_state=42)
model.fit(X_train,y_train)
pred =model.predict(X_test)
acc = f1_score(y_test,pred, average='macro')
acc

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
df

# -- Code Cell --
import numpy as np

# -- Code Cell --
def to_piano_roll(pitches, durations, resolution=256):
    roll = np.zeros((128, resolution))
    t = 0
    for p, d in zip(pitches, durations):
        end = min(int(t + d * resolution), resolution)
        roll[int(p), int(t):end] = 1
        t = end
    return roll

# -- Code Cell --
roll = to_piano_roll(df['pitch'][0], df['duration'][0])
import matplotlib.pyplot as plt
plt.imshow(roll)

# -- Code Cell --
import torch
import torch.nn as nn
from torch.utils.data import DataLoader,Dataset
from torchvision import transforms, models
import torch.optim as optim
from PIL import Image

# -- Code Cell --
class TrainDataset(Dataset):
    def __init__(self, transform):
        self.transform = transform
        self.df = df
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = to_piano_roll(row['pitch'], row['duration'])
        img = torch.tensor(img, dtype=torch.float32)
        label = row['answer']
        if self.transform:
            img = self.transform(img)
        return img, label

# -- Code Cell --
transform = transforms.Compose([
    transforms.Resize((224,224,)),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])

# -- Code Cell --
from torch.utils.data import random_split

# -- Code Cell --
train_ds = TrainDataset(transform=transform)
train_size = int(0.8*len(train_ds))
val_size = len(train_ds) - train_size
train_subset, val_subset = random_split(train_ds,[train_size,val_size])

# -- Code Cell --
val_loader = DataLoader(val_subset, shuffle=False,num_workers=0, batch_size=32 )
train_loader = DataLoader(train_subset, shuffle=True,num_workers=0, batch_size=32 )

# -- Code Cell --
len(df['answer'].unique())

# -- Code Cell --
device =torch.device('cuda')

# -- Code Cell --
class modelmeu(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        in_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Identity()

        self.classifier = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 5)
        )

    def forward(self, x):
        x = self.backbone(x)
        x = self.classifier(x)
        return x

# -- Code Cell --
model = modelmeu()
for param in model.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer=optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max')

# -- Code Cell --
from sklearn.metrics import f1_score

# -- Code Cell --
for epoch in range(30):
    model.train()
    total_loss = 0
    for img,label in train_loader:
        img = img.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss=criterion(output, label)
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for img,label in val_loader:
            img = img.to(device)
            label = label.to(device)
            output = model(img)
            pred = output.argmax(dim=1)
            all_preds.append(pred.cpu())
            all_labels.append(label.cpu())
    all_preds = torch.cat(all_preds)
    all_labels = torch.cat(all_labels)
    all_preds= all_preds.numpy()
    all_labels = all_labels.numpy()
    f1 = f1_score(all_labels,all_preds, average='macro')
    scheduler.step(f1)
    print(f"epoch {epoch+1}, loss{total_loss/len(train_loader)}, f1 {f1}")

# -- Code Cell --
test = pd.read_csv("./test_data.csv")

# -- Code Cell --
df2 = test.groupby("id").agg(list)

# -- Code Cell --
df2

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self, transform):
        self.transform = transform
        self.df = df2
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = to_piano_roll(row['pitch'],row['duration'])
        img = torch.tensor(img, dtype=torch.float32).unsqueeze(0)
        img = img.repeat(3, 1, 1)
        if self.transform:
            img = self.transform(img)
        return img

# -- Code Cell --
test_ds = TestDataset(transform)
test_loader = DataLoader(test_ds, shuffle=False, num_workers=0, batch_size=32)

# -- Code Cell --
preds = []
with torch.no_grad():
    for img in test_loader:
        img = img.to(device)
        output = model(img)
        pred = output.argmax(dim=1)
        preds.append(pred.cpu())
preds_final = torch.cat(preds)
preds_final = preds_final.numpy()
preds_final

# -- Code Cell --
rev_map = {}
for key,value in mapping.items():
    rev_map[value] = key
rev_map

# -- Code Cell --
preds_final_rev = []
for i in range(len(preds_final)):
    preds_final_rev.append(rev_map[preds_final[i]])
preds_final_rev

# -- Code Cell --
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':df2.index.tolist(),
    'answer':preds_final_rev
}).to_csv('subs.csv',index=False)

# -- Markdown Cell --
# # RENUNT LA RESNET

# -- Code Cell --
train["interval"] = train.groupby("id")["pitch"].diff().fillna(0).astype(int)
train["quant_dur"] = (train["duration"] * 4).round() / 4
train["token"] = "i" + train["interval"].astype(str) + "d" + train["quant_dur"].astype(str)
texts = train.groupby("id")["token"].apply(" ".join)
texts

# -- Code Cell --
def get_meta_features(df):
    return df.groupby("id").agg(
        {
            "pitch": ["mean", "std", "min", "max"],
            "duration": ["mean", "std", "sum", "count"],
        }
    )

# -- Code Cell --
'D0.58P29 D0.77P47 D0.14P63 D0.93P23 D0.79P67 D0.93P78 D0.22P28 D0.89P78 D0.14P39 D0.15P75 D0.91P29 D0.45P3'

# -- Code Cell --
text = ""
for duration, pitch in zip(train_grouped['duration'].iloc[2],train_grouped['pitch'].iloc[2]):
    text += "D"
    text += str(round(duration * 4)/4)
    text += "P"
    text += str(pitch)
    text += " "
text

# -- Code Cell --
train

# -- Code Cell --
comp = comp.sort_values('datapointID')
comp

# -- Code Cell --
labels = comp['answer']
labels.unique()

# -- Code Cell --
map_labels = {"Led Zeppelin": 0,'David Bowie':1, 'Psychedelic Corn Cruumpets':2, 'The Beatles':3,'Pantera':4}
lables = labels.map(map_labels)

# -- Code Cell --
lables

# -- Code Cell --
train_grouped = train.groupby('id').agg(list)
texts= []
for i in range(len(train_grouped)):
    text = ""
    for duration, pitch in zip(train_grouped['duration'].iloc[i],train_grouped['pitch'].iloc[i]):
        text += "D"
        text += str(duration)
        text += "P"
        text += str(pitch)
        text += " "
    texts.append(text)

# -- Code Cell --
from sklearn.feature_extraction.text import TfidfVectorizer
tfidf = TfidfVectorizer(ngram_range=(1,2), max_features=3000)
texts_tf = tfidf.fit_transform(texts)

# -- Code Cell --
texts_meta = get_meta_features(train)
texts_meta.columns = ["_".join(col) for col in texts_meta.columns]
texts_meta = texts_meta.loc[texts.index]
import numpy as np

X = np.hstack([texts_tf.toarray(), texts_meta.values])

# -- Code Cell --
X

# -- Code Cell --
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
y = labels
model = RandomForestClassifier(
    n_estimators=400,
    max_features="sqrt",
    n_jobs=-1,
    random_state=42,
    class_weight="balanced",
)
X_train, X_test, y_train, y_test =train_test_split(X,y, test_size=0.2,random_state=42, stratify=y)
model.fit(X_train, y_train)
pred = model.predict(X_test)
acc = f1_score(y_test, pred, average='weighted')
acc

# -- Code Cell --
