import gc
import json
import os
import pandas as pd
import time
import torch

from PIL import Image
from datetime import datetime

from transformers import AutoProcessor, AutoModelForImageTextToText

import use_vlm_ft_gemma3

def clear_memory():
    # Delete variables if they exist in the current global scope
    if "inputs" in globals():
        del globals()["inputs"]
    if "model" in globals():
        del globals()["model"]
    if "processor" in globals():
        del globals()["processor"]
    if "trainer" in globals():
        del globals()["trainer"]
    if "peft_model" in globals():
        del globals()["peft_model"]
    if "bnb_config" in globals():
        del globals()["bnb_config"]
    time.sleep(2)

    # Garbage collection and clearing CUDA memory
    gc.collect()
    time.sleep(2)
    torch.cuda.empty_cache()
    # torch.cuda.synchronize()
    time.sleep(2)
    gc.collect()
    time.sleep(2)

    print(f"GPU allocated memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
    print(f"GPU reserved memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")



if __name__ == "__main__":
    clear_memory()

    in_path_frame_images        = './.../rpp-765k_512'
    in_file_data_entering_test  = './.../test.parquet'

    in_model_name               = "google_gemma-3-4b-pt_local_FT"
    in_name_results_output      = 'results'
    out_path_results            = './.../output'
    os.environ["USED_MODEL"]    = "google_gemma-3-4b_local_FT"

    dict_log = {}
    dict_log['model']           = in_model_name

    print("in_model_name: " + str(in_model_name))

    path_outputs = os.path.join(out_path_results, in_model_name)
    os.makedirs(path_outputs, exist_ok=True)

    df_result = pd.DataFrame(
        columns=['label', 'filename', \
        'brand', 'product_category', 'price', 'regular_price', 'relative_discount', 'absolute_discount', 'GTINs', 'weight_number', 'weight_unit', 'different_types'])
    df_result_cost = pd.DataFrame(
        columns=['label', 'filename']
    )

    df_test = pd.read_parquet(in_file_data_entering_test, engine='pyarrow')
    df_test.reset_index(drop=True, inplace=True)
    
    output_dir = "path-to-checkpoints-directory/google-gemma3-4b-pt/random-subset"
    # Load Model with PEFT adapter
    ft_model = AutoModelForImageTextToText.from_pretrained(
        output_dir,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        attn_implementation="eager",
    )
    processor = AutoProcessor.from_pretrained(output_dir)

    start_index = 0
    output_file = f'{datetime.now().strftime("%Y%m%d_%H%M%S")}_{in_name_results_output}'

    for index, row in df_test.iloc[start_index:].iterrows():
        label       = str(row.label)
        filename    = row.filename

        dict_result = {}
        dict_result['label']    = label
        dict_result['filename'] = filename
        dict_result_cost = {}
        dict_result_cost['label']    = label
        dict_result_cost['filename'] = filename
            
        ############################################
        # PROMPT
        ############################################
        # IMAGE
        image_path = os.path.join( in_path_frame_images, 'test', label, filename )
        pil_image = Image.open(image_path)

        # TASK
        task = "Extract all targets."
        dict_log['prompt_task'] = task

        dict_log, dict_result, dict_result_cost = use_vlm_ft_gemma3.do_request(
            ft_model,
            processor,
            pil_image, 
            task, 
            dict_log,
            dict_result, 
            dict_result_cost,
        )
        
        df_result = pd.concat( [ df_result, pd.DataFrame.from_dict([dict_result]) ], ignore_index=True)
        df_result_cost = pd.concat( [ df_result_cost, pd.DataFrame.from_dict([dict_result_cost]) ], ignore_index=True)

        if index%100 == 0:
            df_result.to_parquet( os.path.join(path_outputs, output_file + '_' + str(index) + '_r.parquet'), index=False, engine='pyarrow')
            df_result_cost.to_parquet( os.path.join(path_outputs, output_file + '_costs_' +  str(index) + '_r.parquet'), index=False, engine='pyarrow')
            df_result = pd.DataFrame( columns=['label', 'filename', \
                        'brand', 'product_category', 'price', 'regular_price', 'relative_discount', 'absolute_discount', 'GTINs', 'weight_number', 'weight_unit', 'different_types'])
            df_result_cost = pd.DataFrame(columns=['label', 'filename'])
        
    
    df_result.to_parquet( os.path.join(path_outputs, output_file + '_' + str(index) + '_r.parquet'), index=False, engine='pyarrow')
    df_result_cost.to_parquet( os.path.join(path_outputs, output_file + '_costs_' +  str(index) + '_r.parquet'), index=False, engine='pyarrow')

    #######################################################################
    #######################################################################
    
    with open(os.path.join(path_outputs, in_name_results_output + '.json'), 'w') as json_file:
        json.dump(dict_log, json_file)