import gradio as gr
import zipfile
from glob import glob

from PIL import Image
from collections import Counter

from huggingface_hub import hf_hub_download, login 
login(token=os.getenv('LOGIN_TOKEN'))
hf_hub_download(repo_id="giniwini/model_creator", filename="ModelCreator.py", )

from ModelCreator import Model

m = Model()

markdown_head = """# <center>🚀👁️ Create your own Classifier 👁️🚀</center>      
### By Elok Quence 🦾

This space is intended to give high level tools so everyone can make his own image classification model and use it for any purpose. 

## How to create it and test it

1. Put some images in a folder classified by subfolders indicating the classes. Like: \n
				- img_folder/cat/image_of_cat_0.png
				- img_folder/dog/image_of_dog_4.jpeg \n
I recomment around 5 images per class. \n
2. Right click in the folder and press "compress..." in ".zip" mode. This will create you a zip file. \n
3. Upload the file on the space. \n
4. Once uploaded, try with some images to test that everything is going well and have fun.  

---
## ⬇️⬇️⬇️ Do you want to export the model or thank my effort? 🫀 Look at the bottom to see the options. ⬇️⬇️⬇️
---
"""

markdown_tail = """
---

**Are you grateful?** 
Consider buying me a coffee subscribing to my [patreon](https://patreon.com/elokquentia?utm_medium=unknown&utm_source=join_link&utm_campaign=creatorshare_creator&utm_content=copyLink) coffee tier \n
or make a donation trough [paypal](https://www.paypal.com/donate/?hosted_button_id=QD2W2G34GWQ4J)
---

## Export the model
Subscribe to patreon Brunch tier to get the password. [patreon](https://patreon.com/elokquentia?utm_medium=unknown&utm_source=join_link&utm_campaign=creatorshare_creator&utm_content=copyLink)

---
## Have you already exported the model? Here's how to use it

### 1. Install the requirements: 
Open the terminal and go to the directory where you placed the requirements.txt file. \n 
Now with this command you can install the dependencies at once. 
```bash
pip install -r requirements.txt
```
Has anything gone wrong? Look inside the requirements.txt and install the dependencies one by one. 

### 2. Use ModelCreator's Model class
Place the ModelCreator.py file on the directory you want to code. 
    - some_folder/ModelCreator.py
    - your_python_script.py

### 3. Use your model
In your_python_script.py, try something like:
```python
from ModelCreator import Model

model_path = 'model.pickle'
m2 = Model().load_model(model_path)

image_path = 'some_image_path/image_1.png'
result = m2(image_path)
print(result)
```
### 4. DO YOU HAVE ANY DOUBTS ON HOW THIS WORKS OR PROBLEMS USING THE MODEL?
Contact me on huggingface or via e-mail genlain@gmail.com
"""


def fit_model(zip_file_path):
    path = 'tmp/'
    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
        zip_ref.extractall(path)

    images_path = glob(f'{path}**/*/*.*')
    labels_train = [i.split('/')[-2] for i in images_path]
    print(labels_train)
    images_train = [Image.open(i) for i in images_path]

    m.train(images_train, labels_train)
    return (f"Model Fitted \n"
            f"Classes: {dict(Counter(labels_train))}")


def predict(image):
    l = m.predict(image)
    return f'Predicted Class: {l}'


def export(password):
    if password == os.getenv('EXPORT_PASSWORD'):
        m.export_model()
        return [[gr.update(visible=True), file] for file in [f'model.pickle', f'requirements.txt', f'ModelCreator.py']]
    else:
        return [None, f'Subscribe to Patreon to Download']*3

with gr.Blocks() as demo:
    gr.Markdown(markdown_head)

    with gr.Row() as g1:
        inp = gr.File(label="zip file")
        out = gr.Textbox()
    btn = gr.Button("Submit Zip File")
    btn.click(fn=fit_model, inputs=inp, outputs=out)

    with gr.Row() as g2:
        inp2 = gr.Image(label="Input Image", type='pil')
        out2 = gr.Textbox()
    btn2 = gr.Button("Predict/Test on an Image")
    btn2.click(fn=predict, inputs=inp2, outputs=out2)

    with gr.Row() as g3:
        inp3 = gr.Textbox()
        out3 = gr.File(label='Download link', visible=True, height=30, interactive=False)
        out4 = gr.File(label='Download link', visible=True, height=30, interactive=False)
        out5 = gr.File(label='Download link', visible=True, height=30, interactive=False)
    btn3 = gr.Button("Export Fitted Model")
    btn3.click(fn=export, inputs=inp3, outputs=[[out3, out3],[out4, out4],[out5, out5]])
    gr.Markdown(markdown_tail)
demo.launch()