import os
import argparse
from torchvision.datasets import EuroSAT
import yaml

def load_config(config_path="config/config.yaml"):
    with open(config_path, "r") as f:
        return yaml.safe_load(f)

def download_eurosat(config):
    eurosat_dir = config["paths"]["eurosat_dir"]
    
    # Check if already downloaded and extracted
    if os.path.exists(os.path.join(eurosat_dir, 'eurosat')) or os.path.exists(os.path.join(eurosat_dir, '2750')):
        print(f"EuroSAT already exists at {eurosat_dir}. Skipping download.")
        return

    print(f"Downloading EuroSAT to {eurosat_dir}...")
    os.makedirs(eurosat_dir, exist_ok=True)
    # EuroSAT dataset from torchvision handles downloading and extracting
    dataset = EuroSAT(root=eurosat_dir, download=True)
    print("EuroSAT download complete!")

def deepglobe_instructions():
    print("-" * 50)
    print("DeepGlobe Dataset Instructions:")
    print("DeepGlobe Land Cover Classification dataset is hosted on Kaggle.")
    print("To download it, please follow these steps:")
    print("1. Ensure you have a Kaggle account and kaggle.json API key configured.")
    print("2. Run the following command in your terminal:")
    print("   kaggle datasets download -d balraj98/deepglobe-land-cover-classification-dataset -p data/DeepGlobe --unzip")
    print("3. Alternatively, download manually from: https://www.kaggle.com/datasets/balraj98/deepglobe-land-cover-classification-dataset")
    print("-" * 50)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download datasets for EcoPulse")
    parser.add_argument("--config", default="config/config.yaml", help="Path to config file")
    args = parser.parse_args()

    config = load_config(args.config)
    
    download_eurosat(config)
    deepglobe_instructions()