import os
import sys
import subprocess
import gdown
import shutil
import nltk
from pathlib import Path

# Install NLTK data
nltk.download('punkt')

# Create directories
os.makedirs('DF-GAN/code/models', exist_ok=True)
os.makedirs('data', exist_ok=True)

# Clone the DF-GAN repository
if not os.path.exists('DF-GAN/.git'):
    print("Cloning DF-GAN repository...")
    subprocess.run(["git", "clone", "https://github.com/tobran/DF-GAN.git", "DF-GAN_temp"])
    
    # Move only necessary files to avoid duplicates
    shutil.copytree('DF-GAN_temp/code/models', 'DF-GAN/code/models', dirs_exist_ok=True)
    shutil.copytree('DF-GAN_temp/code/lib', 'DF-GAN/code/lib', dirs_exist_ok=True)
    
    # Clean up
    shutil.rmtree('DF-GAN_temp')
    
    print("Repository cloned and organized.")

# Download model files
# DF-GAN pretrained bird model
bird_model_url = 'https://drive.google.com/uc?id=1rzfcCvGwU8vLCrn5reWxmrAMms6WQGA6'
bird_model_path = 'data/state_epoch_1220.pth'

# Text encoder for birds
text_encoder_url = 'https://drive.google.com/uc?id=1xwIyLPYtYn9YGPIcRuWXxaxcw_oPGQK4'
text_encoder_path = 'data/text_encoder200.pth'

# Captions DAMSM pickle file
captions_pickle_url = 'https://drive.google.com/uc?id=1FfNMRpOZGaO3mKYyj2VDVEW1ChZ12lJp'
captions_pickle_path = 'data/captions_DAMSM.pickle'

# Download if files don't exist
if not os.path.exists(bird_model_path):
    print(f"Downloading bird model to {bird_model_path}...")
    gdown.download(bird_model_url, bird_model_path, quiet=False)

if not os.path.exists(text_encoder_path):
    print(f"Downloading text encoder to {text_encoder_path}...")
    gdown.download(text_encoder_url, text_encoder_path, quiet=False)

if not os.path.exists(captions_pickle_path):
    print(f"Downloading captions pickle to {captions_pickle_path}...")
    gdown.download(captions_pickle_url, captions_pickle_path, quiet=False)

print("All model files downloaded and prepared successfully!")