os.environ["CUDA_VISIBLE_DEVICES"]="3"

if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available. Using GPU.")
else:
    device = torch.device("cpu")
    print("CUDA is not available. Using CPU.")


registry.load_database('finetune.yml')
protocol = registry.get_protocol("dis.SpeakerDiarization.aa", preprocessors={"audio":
                                                        FileFinder()})

from pyannote.audio import Model
model = Model.from_pretrained("pyannote/segmentation-3.0", use_auth_token="token_here")

from pyannote.audio.tasks import Segmentation
task = Segmentation(
    protocol, 
    duration=model.specifications.duration, 
    max_num_speakers=len(model.specifications.classes), 
    batch_size=32,
    num_workers=0, 
    loss="bce", 
    vad_loss="bce")
model.task = task
model.prepare_data()
model.setup()

model = model.to(device)

save_dir = 'save_folder'
# this takes approximately 15min to run on Google Colab GPU
from types import MethodType
from torch.optim import Adam
from pytorch_lightning.callbacks import (
    # EarlyStopping,
    ModelCheckpoint,
    RichProgressBar,
)

# we use Adam optimizer with 1e-4 learning rate
def configure_optimizers(self):
    return Adam(self.parameters(), lr=1e-4)

model.configure_optimizers = MethodType(configure_optimizers, model)

monitor, direction = task.val_monitor
checkpoint = ModelCheckpoint(
    monitor=monitor,
    mode=direction,
    # save_top_k=1,
    every_n_epochs=1,
    save_last=True,
    save_weights_only=False,
    filename="{epoch}",
    verbose=False,

callbacks = [RichProgressBar(), checkpoint]

from pytorch_lightning import Trainer

trainer = Trainer(accelerator="gpu", 
                  callbacks=callbacks, 
                  max_epochs=20,
                  limit_val_batches=0,
                  num_sanity_val_steps=0,
                  logger = False)
trainer.fit(model)

print('done training')

base_fine = 'last.ckpt'

from pyannote.audio.pipelines import SpeakerDiarization as SpeakerDiarizationPipeline

STEP=3.0

# hyperparameters used for AMI
PIPELINE_PARAMS = {
    "clustering": {
        "method": "centroid",
        "min_cluster_size": 15,
        "threshold": 0.6285824248662424,
    },
    "segmentation": {
        "min_duration_off": 0.0,
    },
}
print('setting up model')
adapted_model = Model.from_pretrained(base_fine)
pipeline_adapted = SpeakerDiarizationPipeline(adapted_model, embedding="speechbrain/spkrec-ecapa-voxceleb", clustering="AgglomerativeClustering").to(device)
pipeline_adapted.instantiate(PIPELINE_PARAMS)
print('done setting')

for file in protocol.test() : 
    fine_save = 'save_results_folder'

    uri = file['uri']
    
    fine_save = os.path.join(fine,f'{ uri }.rttm')

    print('getting output')
    out_adapted = pipeline_adapted(file)
    print('got it')
    print(f"writting fine tune rttm for {uri}")
    with open(fine_save,'w') as fine_rttm : 
        out_adapted.write_rttm(fine_rttm)
    print("done")