import numpy as np
import torch
import torchvision.transforms as T
from decord import VideoReader, cpu
from PIL import Image
from torchvision.transforms.functional import InterpolationMode
from transformers import AutoModel, AutoTokenizer

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform


def load_image(image_file):
    transform = build_transform(input_size=800)
    pixel_values = transform(image_file).unsqueeze(0)  # Add batch dimension
    return pixel_values


def main(image_path,model,tokenizer):

    pixel_values = load_image(image_path).to(torch.float32).to("cpu")
    generation_config = dict(max_new_tokens=1024, do_sample=True)
    
    question = """<image>\n**Instruction:**
    Analyze the image to extract values for the specified keys. Use the detailed descriptions below to determine the correct value for each key. Handle missing or ambiguous data as instructed.
    
    ---
    ### Keys and Descriptions
    
    1. **`Surat Ketetapan Kewajiban Pembayaran PKB/BBNKB,SWDKLLJ DAN PNBP`**
       - **Extract**: The value of the field is"Surat Ketetapan Kewajiban Pembayaran PKB/BBNKB, SWDKLLJ DAN PNBP" and this is title.
       - **If the Field is Present** : "present"
       - **If the Field is Absent**: `"null"`
    
    2. **`nomor_registrasi`**
       - **Extract**: The "NOMOR REGISTRASI" field.
       - **If the Field is Absent**: `"null"`
       - **If the Field is Present but No Value is Provided**: `"empty"`
    
    3.  **`nama_pemilik`**
       - **Extract**: The "NAMA PEMILIK" field.
       - **If the Field is Absent**: `"null"`
       - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    4. **`alamat`**
       - **Extract**: The "ALAMAT" field and this is address may contain 1 or more lines format : "BUNTET PESANTREN RT/RW/015/005 DESA MERTAPADA
    KULON KEC. ASTANAJAPURA KAB.CIREBON ".
       - **If the Field is Absent**: `"null"`
       - **If the Field is Present but No Value is Provided**: `"empty"`
    
    5. **`merk/type`**
       - **Extract**: The "MERK/TYPE" field.
       - **If the Field is Absent**: `"null"`
       - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    6. **`jenis/model`**
       - **Extract**: The "JENIS/MODEL" field.
       - **If the Field is Absent**: `"null"`
       - **If the Field is Present but No Value is Provided**: `"empty"`
    
    7.  **`tahun_pembuatan`**
        - **Extract**: The "TAHUN PEMBUATAN/PERAKITAN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    8.  **`warna_kb`**
        - **Extract**: The "WARNA KB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    9. **`isi_silinder`**
        - **Extract**: The "ISI SILINDER/HP" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    10. **`nomor_rangka`**
        - **Extract**: The "NOMOR RANGKA/NIK" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    11.**`nomor_mesin`**
        - **Extract**: The "NOMOR MESIN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    12. **`no_bpkp`**
        - **Extract**: The "NO BPKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    13. **`berlaku s/d`**
        - **Extract**: The "BERLAKU S/D" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    14. **`bahan_bakar`**
        - **Extract**: The "BAHAN BAKAR" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    15. **`warna_tnkb`**
        - **Extract**: The "WARNA TNKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    
    16.  **`kepemilikan_ke`**
        - **Extract**: The "KEPEMILIKAN KE" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    17. **`no_registrasi_lama`**
        - **Extract**: The "NO REGISTRASI LAMA" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    18. **`kode_njkb`**
        - **Extract**: The "KODE NJKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    19. **`no`**
        - **Extract**: The "NO." field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    20.**`asal_daerah`**
        - **Extract**: The "ASAL DAERAH" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    21. **`no_urut`**
        - **Extract**: The "NO. URUT" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    22. **`no_kohir`**
        - **Extract**: The "NO. KOHIR" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    23. **`nik_no.hp`**
        - **Extract**: The "NIK/NO.HP" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    24. **`pkok_bbnkb`**
        - **Extract**: The "PKOK BBNKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    25. **`pkok_pkb`**
        - **Extract**: The "POKOK PKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    26. **`pkok_swdkljj`**
        - **Extract**: The "POKOK SWDKLJJ" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    27. **`pkok_penerbitan_stnk`**
        - **Extract**: The "POKOK PENERBITAN STNK" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    28. **`pkok_penerbitan_tnkb`**
        - **Extract**: The "POKOK PENERBITAN TNKB/NRKB PILIHAN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    29. **`pkok_jumlah`**
        - **Extract**: The "POKOK JUMLAH" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    30. **`sanksi_adm_bbnkb`**
        - **Extract**: The "SANKSI ADM BBNKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    31. **`sanksi_adm_pkb`**
        - **Extract**: The "SANKSI ADM PKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    32. **`sanksi_adm_swdkljj`**
        - **Extract**: The "SANKSI ADM SWDKLJJ" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    33. **`sanksi_adm_penerbitan_stnk`**
        - **Extract**: The "SANKSI ADM PENERBITAN STNK" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    34. **`sanksi_adm_penerbitan_tnkb`**
        - **Extract**: The "SANKSI ADM PENERBITAN TNKB/NRKB PILIHAN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    35. **`sanksi_adm_jumlah`**
        - **Extract**: The "SANKSI ADM JUMLAH " field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    36. **`jumlah_bbnkb`**
        - **Extract**: The "JUMLAH BBNKB" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    37. **`jumlah_pkb`**
        - **Extract**: The "JUMLAH PKB " field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    38. **`jumlah_swdkljj`**
        - **Extract**: The "JUMLAH SWDKLJJ" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    39. **`jumlah_penerbitan_stnk`**
        - **Extract**: The "JUMLAH PENERBITAN STNK" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    40. **`jumlah_penerbitan_tnkb`**
        - **Extract**: The "JUMLAH PENERBITAN TNKB/NRKB PILIHAN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    41. **`jumlah_jumlah`**
        - **Extract**: The "JUMLAH JUMLAH" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    42. **`ditetapkan_tanggal`**
        - **Extract**: The "DITETAPKAN TANGGAL" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    43. **`penetapan`**
        - **Extract**: The "PENETAPAN" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    44. **`kasir`**
        - **Extract**: The "KASIR" field.
        - **If the Field is Absent**: `"null"`
        - **If the Field is Present but No Value is Provided**: `"empty"`
    
    ---
    ### Reasoning Process
    For each key, explain your reasoning:
    Indicate whether the field was present.
    Justify the extracted value or the use of "null" or "empty" as per the conditions.
    
    ### Output Format
    
    ```json
    {
      "Surat_Ketetapan_Kewajiban_Pembayaran_PKB_BBNKB_SWDKLLJ_DAN_PNBP": "present OR null",
      "nomor_registrasi": "<value> OR empty OR null",
      "nama_pemilik": "<value> OR empty OR null",
      "alamat": "<value> OR empty OR null",
      "merk_type": "<value> OR empty OR null",
      "jenis_model": "<value> OR empty OR null",
      "tahun_pembuatan": "<value> OR empty OR null",
      "warna_kb": "<value> OR empty OR null",
      "isi_silinder": "<value> OR empty OR null",
      "nomor_rangka": "<value> OR empty OR null",
      "nomor_mesin": "<value> OR empty OR null",
      "no_bpkp": "<value> OR empty OR null",
      "berlaku_s_d": "<value> OR empty OR null",
      "bahan_bakar": "<value> OR empty OR null",
      "warna_tnkb": "<value> OR empty OR null",
      "kepemilikan_ke": "<value> OR empty OR null",
      "no_registrasi_lama": "<value> OR empty OR null",
      "kode_njkb": "<value> OR empty OR null",
      "no": "<value> OR empty OR null",
      "asal_daerah": "<value> OR empty OR null",
      "no_urut": "<value> OR empty OR null",
      "no_kohir": "<value> OR empty OR null",
      "nik_no_hp": "<value> OR empty OR null",
      "pkok_bbnkb": "<value> OR empty OR null",
      "pkok_pkb": "<value> OR empty OR null",
      "pkok_swdkljj": "<value> OR empty OR null",
      "pkok_penerbitan_stnk": "<value> OR empty OR null",
      "pkok_penerbitan_tnkb": "<value> OR empty OR null",
      "pkok_jumlah": "<value> OR empty OR null",
      "sanksi_adm_bbnkb": "<value> OR empty OR null",
      "sanksi_adm_pkb": "<value> OR empty OR null",
      "sanksi_adm_swdkljj": "<value> OR empty OR null",
      "sanksi_adm_penerbitan_stnk": "<value> OR empty OR null",
      "sanksi_adm_penerbitan_tnkb": "<value> OR empty OR null",
      "sanksi_adm_jumlah": "<value> OR empty OR null",
      "jumlah_bbnkb": "<value> OR empty OR null",
      "jumlah_pkb": "<value> OR empty OR null",
      "jumlah_swdkljj": "<value> OR empty OR null",
      "jumlah_penerbitan_stnk": "<value> OR empty OR null",
      "jumlah_penerbitan_tnkb": "<value> OR empty OR null",
      "jumlah_jumlah": "<value> OR empty OR null",
      "ditetapkan_tanggal": "<value> OR empty OR null",
      "penetapan": "<value> OR empty OR null",
      "kasir": "<value> OR empty OR null"
    }
    }
    Return Output:
    Generate a JSON object:
    {
        "reasoning": "reasoning for each key",
        "output JSON": "key-value pairs"
    }
    —
    
    """
    response = model.chat(tokenizer, pixel_values, question, generation_config)
    return (f'User: {question}\nAssistant: {response}')