#!/usr/bin/env python3
"""
Script to upload PhaseNet-TF model to Hugging Face Hub
"""
import os
import json
from pathlib import Path
from huggingface_hub import HfApi, create_repo, upload_file
from huggingface_hub import hf_hub_download
import torch
import yaml

def create_model_card(version=None):
    """Create a comprehensive model card for PhaseNet-TF"""
    title = f"# PhaseNet-TF Alaska"
    if version:
        title = f"# PhaseNet-TF Alaska - {version}"
    
    return f"""{title}

## Model Description

PhaseNet-TF is an advanced deep learning model for automatic seismic phase picking (P-wave, S-wave, and PS-wave detection) using spectrogram-based image segmentation approaches. The model leverages DeepLabV3Plus architecture to detect seismic arrivals with high accuracy, especially for weak and noisy signals from ocean-bottom seismometers and weak phases such as slab interface refracted PS and SP waves. This Alaska version is specifically trained on the PS_Alaska dataset for P and S phases. For more details, please refer to the paper and the [PhaseNet-TF](https://github.com/swei-seismo/PhaseNet-TF) repository.

## Model Architecture

- **Backbone**: DeepLabV3Plus with ResNet34 encoder
- **Input**: 3-component seismic waveforms converted to 6-channel spectrograms (real + imaginary)
- **Output**: Probability maps for P, S, PS phases and noise
- **Sampling Rate**: 40 Hz (dt_s = 0.025s)
- **Window Length**: 4800 points (120 seconds)
- **Spectrogram Size**: 64 × 4800 (frequency × time)
- **Input Channels**: 6 (3 real + 3 imaginary spectrogram channels)
- **Output Classes**: 4 (noise, P, S, PS)
"""

def create_config_json(model_path, version=None):
    """Create config.json with model metadata"""
    config = {
        "model_type": "phasenet-tf",
        "architecture": "DeepLabV3Plus with ResNet34 encoder",
        "input_channels": 6,  # 3-component real + 3-component imaginary spectrograms
        "output_classes": 4,  # noise, P, S, PS
        "sampling_rate": 40,  # 1/0.025 = 40 Hz
        "window_length": 4800,  # 120 seconds at 40 Hz
        "phases": ["P", "S", "PS"],
        "framework": "pytorch",
        "license": "mit",
        "tags": ["seismic", "earthquake", "phase-picking", "deep-learning", "deeplabv3plus"]
    }
    
    if version:
        config["version"] = version
        config["checkpoint_file"] = f"alaska_{version}.bin"
    
    return config

def create_main_readme():
    """Create a main README that showcases both versions"""
    return """---
language: en
tags:
- seismic
- earthquake
- phase-picking
- deep-learning
- pytorch
license: mit
datasets:
- PS_Alaska
metrics:
- f1-score
- precision
- recall
---

# PhaseNet-TF Alaska: Advanced Seismic Arrival Time Detection

## Model Description

PhaseNet-TF is an advanced deep learning model for automatic seismic phase picking (P-wave, S-wave, and PS-wave detection) using spectrogram-based image segmentation approaches. The model leverages DeepLabV3Plus architecture to detect seismic arrivals with high accuracy, especially for weak and noisy signals from ocean-bottom seismometers and weak phases such as slab interface refracted PS and SP waves. This Alaska version is specifically trained on the PS_Alaska dataset for P and S phases.

## Available Versions

This repository contains two versions of the PhaseNet-TF Alaska model:

### 🔄 Iteration 1
- **Model File**: `alaska_iter1.bin`
- **Config**: `config_iter1.json`
- **Documentation**: [README_iter1.md](README_iter1.md)

### 🔄 Iteration 2  
- **Model File**: `alaska_iter2.bin`
- **Config**: `config_iter2.json`
- **Documentation**: [README_iter2.md](README_iter2.md)

## Model Architecture

- **Backbone**: DeepLabV3Plus with ResNet34 encoder
- **Input**: 3-component seismic waveforms converted to 6-channel spectrograms (real + imaginary)
- **Output**: Probability maps for P, S, PS phases and noise
- **Sampling Rate**: 40 Hz (dt_s = 0.025s)
- **Window Length**: 4800 points (120 seconds)
- **Spectrogram Size**: 64 × 4800 (frequency × time)
- **Input Channels**: 6 (3 real + 3 imaginary spectrogram channels)
- **Output Classes**: 4 (noise, P, S, PS)

## Citation

If you use this model in your research, please cite:

```bibtex
@article{jie2025background,
  title={Background Seismicity and Aftershocks of the 2020-2021 Large Earthquakes at the Alaska Peninsula Revealed by a Deep-learning-based Catalog},
  author={Jie, Yaqi and Wei, Songqiao Shawn and Zhu, Weiqiang and Freymueller, Jeffrey Todd and Elliott, Julie},
  journal={Authorea Preprints},
  year={2025},
  publisher={Authorea}
}
```

## License

This model is licensed under the MIT License.
"""

def upload_model_to_hf(
    checkpoint_path: str,
    config_path: str = None,
    repo_name: str = "PhaseNet-TF_Alaska",
    username: str = None,
    token: str = None,
    version: str = None,
    upload_main_readme: bool = False
):
    """Upload model to Hugging Face Hub"""
    
    # Initialize API
    if token:
        api = HfApi(token=token)
    else:
        api = HfApi()
    
    # Get username if not provided
    if username is None:
        try:
            username = api.whoami()["name"]
            print(f"Using logged-in username: {username}")
        except Exception as e:
            print(f"Error getting username: {e}")
            print("Please provide username with --username or login with huggingface-cli login")
            return
    
    # Create repository
    repo_id = f"{username}/{repo_name}"
    try:
        if token:
            create_repo(repo_id, token=token, exist_ok=True)
        else:
            create_repo(repo_id, exist_ok=True)
        print(f"Repository {repo_id} created/accessed successfully")
    except Exception as e:
        print(f"Error creating repository: {e}")
        return
    
    # Determine file paths based on version
    if version:
        model_path = f"alaska_{version}.bin"
        config_path_in_repo = f"config_{version}.json"
        readme_path_in_repo = f"README_{version}.md"
        print(f"Uploading version {version} to {repo_id}")
    else:
        model_path = "alaska.bin"
        config_path_in_repo = "config.json"
        readme_path_in_repo = "README.md"
        print(f"Uploading to {repo_id}")
    
    # Upload checkpoint
    print(f"Uploading model checkpoint as {model_path}...")
    upload_file(
        path_or_fileobj=checkpoint_path,
        path_in_repo=model_path,
        repo_id=repo_id,
        token=token
    )
    
    # Upload config if provided
    if config_path and os.path.exists(config_path):
        config_yaml_path = f"config_{version}.yaml" if version else "config.yaml"
        print(f"Uploading config file as {config_yaml_path}...")
        upload_file(
            path_or_fileobj=config_path,
            path_in_repo=config_yaml_path,
            repo_id=repo_id,
            token=token
        )
    
    # Create and upload config.json
    config_json = create_config_json(checkpoint_path, version)
    config_json_path = "config.json"
    with open(config_json_path, 'w') as f:
        json.dump(config_json, f, indent=2)
    
    upload_file(
        path_or_fileobj=config_json_path,
        path_in_repo=config_path_in_repo,
        repo_id=repo_id,
        token=token
    )
    
    # Create and upload README.md
    model_card = create_model_card(version)
    
    model_card += """

## Citation

If you use this model in your research, please cite:

```bibtex
@article{jie2025background,
  title={Background Seismicity and Aftershocks of the 2020-2021 Large Earthquakes at the Alaska Peninsula Revealed by a Deep-learning-based Catalog},
  author={Jie, Yaqi and Wei, Songqiao Shawn and Zhu, Weiqiang and Freymueller, Jeffrey Todd and Elliott, Julie},
  journal={Authorea Preprints},
  year={2025},
  publisher={Authorea}
}
```

## License

This model is licensed under the MIT License.
"""
    
    readme_path = "README.md"
    with open(readme_path, 'w') as f:
        f.write(model_card)
    
    upload_file(
        path_or_fileobj=readme_path,
        path_in_repo=readme_path_in_repo,
        repo_id=repo_id,
        token=token
    )
    
    # Upload main README if requested
    if upload_main_readme:
        print("Uploading main README.md...")
        main_readme = create_main_readme()
        main_readme_path = "main_README.md"
        with open(main_readme_path, 'w') as f:
            f.write(main_readme)
        
        upload_file(
            path_or_fileobj=main_readme_path,
            path_in_repo="README.md",
            repo_id=repo_id,
            token=token
        )
        os.remove(main_readme_path)
    
    # Clean up temporary files
    os.remove(config_json_path)
    os.remove(readme_path)
    
    print(f"Model uploaded successfully to https://huggingface.co/{repo_id}")
    if version:
        print(f"Files uploaded: {model_path}, {config_path_in_repo}, {readme_path_in_repo}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Upload PhaseNet-TF model to Hugging Face")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint (.ckpt)")
    parser.add_argument("--config", help="Path to config file (.yaml)")
    parser.add_argument("--repo-name", default="PhaseNet-TF_Alaska", help="Repository name")
    parser.add_argument("--username", help="Hugging Face username (optional if already logged in)")
    parser.add_argument("--token", help="Hugging Face token (optional if already logged in)")
    parser.add_argument("--version", help="Version suffix for file naming (e.g., iter1, iter2, v1, v2)")
    parser.add_argument("--main-readme", action="store_true", help="Upload main README that showcases both versions")
    
    args = parser.parse_args()
    
    upload_model_to_hf(
        checkpoint_path=args.checkpoint,
        config_path=args.config,
        repo_name=args.repo_name,
        username=args.username,
        token=args.token,
        version=args.version,
        upload_main_readme=args.main_readme
    )