### This is example of the script that will be run in the test environment.
### Some parts of the code are compulsory and you should NOT CHANGE THEM.
### They are between '''---compulsory---''' comments.
### You can change the rest of the code to define and test your solution.
### However, you should not change the signature of the provided function.
### The script would save "submission.parquet" file in the current directory.
### The actual logic of the solution is implemented in the `handcrafted_solution.py` file.
### The `handcrafted_solution.py` file is a placeholder for your solution.
### You should implement the logic of your solution in that file.
### You can use any additional files and subdirectories to organize your code.

'''---compulsory---'''
import hoho; hoho.setup() # YOU MUST CALL hoho.setup() BEFORE ANYTHING ELSE
import subprocess
import importlib
from pathlib import Path
import subprocess


### The function below is useful for installing additional python wheels.        
def install_package_from_local_file(package_name, folder='packages'):
    """
    Installs a package from a local .whl file or a directory containing .whl files using pip.

    Parameters:
    path_to_file_or_directory (str): The path to the .whl file or the directory containing .whl files.
    """
    try:
        pth = str(Path(folder) / package_name)
        subprocess.check_call([subprocess.sys.executable, "-m", "pip", "install", 
                               "--no-index",  # Do not use package index
                               "--find-links", pth,  # Look for packages in the specified directory or at the file
                               package_name])  # Specify the package to install
        print(f"Package installed successfully from {pth}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to install package from {pth}. Error: {e}")
        

# pip download webdataset -d packages/webdataset --platform manylinux1_x86_64 --python-version 38 --only-binary=:all:
install_package_from_local_file('webdataset')
# install_package_from_local_file('tqdm')

### Here you can import any library or module you want.
### The code below is used to read and parse the input dataset.
### Please, do not modify it.

import webdataset as wds
from tqdm import tqdm
from typing import Dict
import pandas as pd
from transformers import AutoTokenizer
import os
import time
import io
from read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
from PIL import Image as PImage
import numpy as np

    
def proc(row, split='train'):
    out = {}
    for k, v in row.items():
        colname = k.split('.')[0]
        if colname in {'ade20k', 'depthcm', 'gestalt'}:
            if colname in out:
                out[colname].append(v)
            else:
                out[colname] = [v]
        elif colname in {'wireframe', 'mesh'}:
            # out.update({a: b.tolist() for a,b in v.items()})
            out.update({a: b for a,b in v.items()})
        elif colname in 'kr':
            out[colname.upper()] = v
        else:
            out[colname] = v
            
    return Sample(out)


class Sample(Dict):
    def __repr__(self):
        return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k,v in self.items()})

def convert_entry_to_human_readable(entry):
    out = {}
    already_good = ['__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces', 'face_semantics', 'K', 'R', 't']
    for k, v in entry.items():
        if k in already_good:
            out[k] = v
            continue
        if k == 'points3d':
            out[k] = read_points3D_binary(fid=io.BytesIO(v))
        if k == 'cameras':
            out[k] = read_cameras_binary(fid=io.BytesIO(v))
        if k == 'images':
            out[k] = read_images_binary(fid=io.BytesIO(v))
        if k in ['ade20k', 'gestalt']:
            out[k] =  [PImage.open(io.BytesIO(x)).convert('RGB') for x in v]
        if k == 'depthcm':
            out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']]
    return out

'''---end of compulsory---'''

def download_package(package_name, path_to_save='packages'):
    """
    Downloads a package using pip and saves it to a specified directory.

    Parameters:
    package_name (str): The name of the package to download.
    path_to_save (str): The path to the directory where the package will be saved.
    """
    try:
        # pip download webdataset -d packages/webdataset --platform manylinux1_x86_64 --python-version 38 --only-binary=:all:
        subprocess.check_call([subprocess.sys.executable, "-m", "pip", "download", package_name, 
                               "-d", str(Path(path_to_save)/package_name),  # Download the package to the specified directory
                               "--platform", "manylinux1_x86_64",  # Specify the platform
                               "--python-version", "38",  # Specify the Python version
                               "--only-binary=:all:"])  # Download only binary packages
        print(f'Package "{package_name}" downloaded successfully')
    except subprocess.CalledProcessError as e:
        print(f'Failed to downloaded package "{package_name}". Error: {e}')


### The part below is used to define and test your solution.

if __name__ == "__main__":
    from handcrafted_solution import predict
    print ("------------ Loading dataset------------ ")
    params = hoho.get_params()
    dataset = hoho.get_dataset(decode=None, split='all', dataset_type='webdataset')
    print('------------ Now you can do your solution ---------------')
    solution = []
    for i, sample in enumerate(tqdm(dataset)):
        pred_vertices, pred_edges, semantics = predict(sample, visualize=False)
        solution.append({
                        '__key__': sample['__key__'], 
                        'wf_vertices': pred_vertices.tolist(),
                        'wf_edges': pred_edges,
                        'edge_semantics': semantics,
                    })
    print('------------ Saving results ---------------')
    sub = pd.DataFrame(solution, columns=["__key__", "wf_vertices", "wf_edges", "edge_semantics"])
    sub.to_parquet(Path(params['output_path']) / "submission.parquet")
    print("------------ Done ------------ ")