#!/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():
    """Create a comprehensive model card for PhaseNet-TF"""
    return """---
language: en
tags:
- seismic
- earthquake
- phase-picking
- deep-learning
- pytorch
license: mit
datasets:
- PS_Alaska
metrics:
- f1-score
- precision
- recall
---

# PhaseNet-TF: 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. 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)

## Load the checkpoint
checkpoint = torch.load("alaska_iter2.ckpt", map_location="cpu")

## 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 create_config_json(model_path):
    """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"]
    }
    return config

def upload_model_to_hf(
    checkpoint_path: str,
    config_path: str = None,
    repo_name: str = "PhaseNet-TF_Alaska",
    username: str = None,
    token: str = None
):
    """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
    
    # Upload checkpoint
    print("Uploading model checkpoint...")
    upload_file(
        path_or_fileobj=checkpoint_path,
        path_in_repo="pytorch_model.bin",
        repo_id=repo_id,
        token=token
    )
    
    # Upload config
    if config_path and os.path.exists(config_path):
        print("Uploading config file...")
        upload_file(
            path_or_fileobj=config_path,
            path_in_repo="config.yaml",
            repo_id=repo_id,
            token=token
        )
    
    # Create and upload config.json
    config_json = create_config_json(checkpoint_path)
    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.json",
        repo_id=repo_id,
        token=token
    )
    
    # Create and upload README.md
    model_card = create_model_card()
    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.md",
        repo_id=repo_id,
        token=token
    )
    
    # 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 __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)")
    
    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
    )