from huggingface_hub import hf_hub_download, Repository
import gym
import numpy as np
import os

# Define your username and repo name
username = "willco-afk"  # Your Hugging Face username
repo_name = "frozenslippery"  # Your Hugging Face Space name

# Correct file path where the Q-table is located
repo_id = "willco-afk/frozenslippery"
file_path = "q_table_frozenlake.npy"  # Path to the Q-table file in the repo

# Try downloading the Q-table
try:
    download_path = hf_hub_download(repo_id=repo_id, filename=file_path)
    # Load the Q-table
    q_table = np.load(download_path)
except Exception as e:
    print(f"Error downloading the Q-table: {e}")
    # Handle the error (for example, by uploading the Q-table manually if needed)

# Check if the Q-table was loaded successfully
if 'q_table' in locals():
    # Save the model (Q-table) as a .npz file in the repo's folder
    model_filename = "q_table_frozenlake.npz"
    np.save(model_filename, q_table)

    # Initialize the Hugging Face repo for the Space (no need to create it again)
    repo = Repository(local_dir=repo_name, clone_from=f"{username}/{repo_name}")

    # Add and push the model file to Hugging Face Hub
    repo.git_add(model_filename)  # Add the Q-table to the repo
    repo.git_commit("Add trained Q-table")  # Commit the Q-table
    repo.git_push()  # Push the changes to Hugging Face Hub

    # Write the README file with details
    readme_content = "# FrozenLake RL Model\n\n"
    readme_content += "This model represents a Q-learning agent for the `FrozenLake-v1` environment with `is_slippery=True`.\n\n"
    readme_content += "### Usage Instructions\n\n"
    readme_content += "To use this model, you need to initialize the FrozenLake environment using OpenAI's gym:\n\n"
    readme_content += "```python\n"
    readme_content += "import gym\n"
    readme_content += "env = gym.make('FrozenLake-v1', is_slippery=True)\n"
    readme_content += "```\n\n"
    readme_content += "### Model Details\n\n"
    readme_content += "This model uses a Q-table learned through Q-learning in the `FrozenLake-v1` environment. The agent was trained using the following parameters:\n\n"
    readme_content += "- **Learning Rate:** 0.1\n"
    readme_content += "- **Discount Factor (gamma):** 0.99\n"
    readme_content += "- **Exploration Rate (epsilon):** Decays from 1.0 to 0.01\n"
    readme_content += "- **Training Episodes:** 1000\n"
    readme_content += "- **Max Steps per Episode:** 100\n\n"
    readme_content += "### About the Environment\n\n"
    readme_content += "The `FrozenLake-v1` environment is a gridworld where the agent must navigate a frozen lake while avoiding holes. It can slip based on the `is_slippery` parameter, making the environment stochastic.\n"

    # Write the README file
    with open(f"{repo_name}/README.md", "w") as readme_file:
        readme_file.write(readme_content)

    # Add and push the README file
    repo.git_add("README.md")
    repo.git_commit("Add README for FrozenLake RL model")
    repo.git_push()

else:
    print("Q-table was not loaded, skipping further operations.")