from typing import Dict, Optional, List, Literal
from dataclasses import dataclass, field
import json
import transformers

from utils.supported_models import MODEL_HF_PATH, MODEL_FAMILIES


@dataclass
class ModelArguments:
    model_id: str = field(default="llava-1.5-7b")
    model_local_path: Optional[str] = field(default=None)
    cut_embed_dim: Optional[int] = field(
        default=3072,
        metadata={"help": "cut off embed dim for tbstars model"},
    )
    parscale_n: Optional[int] = field(
        default=2,
        metadata={"help": "Use knn for partial fc"},
    )
    parscale_n_tokens: Optional[int] = field(
        default=42,
        metadata={"help": "parscale_n_tokens"},
    )
    parscale_attn_smooth: Optional[float] = field(
        default=0.01,
        metadata={"help": "parscale_attn_smooth"},
    )
    parscale_layer: Optional[int] = field(
        default=1,
        metadata={"help": "when to begin the parscale"},
    )
    perceiver_num_heads: Optional[int] = field(
        default=32,
        metadata={"help": "qformer_head"},
    )
    perceiver_num_queries: Optional[int] = field(
        default=144,
        metadata={"help": "qformer_dim"},
    )
    lora_load: Optional[bool] = field(
        default=False,
        metadata={"help": "if loading the lora pth"},
    )
    lora_pth: Optional[str] = field(
        default='',
        metadata={"help": "lora pth"},
    )



    def __post_init__(self):
        assert self.model_id in MODEL_HF_PATH, f"Unknown model_id: {self.model_id}"
        self.model_hf_path: str = MODEL_HF_PATH[self.model_id]
        assert self.model_id in MODEL_FAMILIES, f"Unknown model_id: {self.model_id}"
        self.model_family_id: str = MODEL_FAMILIES[self.model_id]

        if not self.model_local_path:
            self.model_local_path = self.model_hf_path


@dataclass
class DataArguments:
    data_path: str = field(
        default=None, metadata={"help": "Path to the training data json file."}
    )
    eval_data_path: Optional[str] = field(
        default=None, metadata={"help": "Path to the evaluation data json file."}
    )
    image_folder: Optional[str] = field(default=None)
    video_folder: Optional[str] = field(default=None)
    num_frames: Optional[int] = field(default=8)
    user_key: Optional[str] = field(default="human")
    assistant_key: Optional[str] = field(default="gpt")
    image_data_path: str = field(
        default=None, metadata={"help": "Path to the image data json file."}
    )
    text_data_path: str = field(
        default=None, metadata={"help": "Path to the text data json file."}
    )
    query_data_path: str = field(
        default=None, metadata={"help": "Path to the query data json file."}
    )
    cand_pool_path: str = field(
        default=None, metadata={"help": "Path to the cand pool data json file."}
    )
    instructions_path: str = field(
        default=None, metadata={"help": "Path to the instructions data json file."}
    )
    rerank_data_path: str = field(
        default=None, metadata={"help": "Path to the rerank data json file."}
    )
    image_path_prefix: str = field(
        default=None, metadata={"help": "Path to the image files."}
    )
    tables: str = field(
        default=None, metadata={"help": "Odps Tables."}
    )
    tables_eval: str = field(
        default=None, metadata={"help": "Odps Tables for eval."}
    )
    max_lines: int = field(
        default=-1, metadata={"help": "Max data lines."}
    )
    max_lines_eval: int = field(
        default=-1, metadata={"help": "Max data lines of eval."}
    )
    image_type: str = field(
        default="crop", metadata={"help": "eval instruction."}
    )
    text_require: Optional[bool] = field(
        default=False,
        metadata={"help": "requrei text or not"},
    )
    


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    model_max_length: int = field(
        default=1024,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    eval_json: str = field(
        default='eval_results.jsonl', metadata={"help": "json file to save eval result"}
    )

    sub_loss: Optional[bool] = field(
        default=True,
        metadata={"help": "if use subspace loss"},
    )
    sub_loss_weight: Optional[float] = field(
        default=1,
        metadata={"help": "sub loss weight"},
    )
    margin_value: Optional[float] = field(
        default=1,
        metadata={"help": "margin value of triplet loss"},
    )
    triplet_weight: Optional[float] = field(
        default=1,
        metadata={"help": "margin value of triplet loss"},
    )
    club_lr: float = field(
        default=1e-5, metadata={"help": "lambda for uniform loss"}
    )
    club_weight_decay: float = field(
        default=0.01, metadata={"help": "lambda for uniform loss"}
    )
    club_flag: bool = field(
        default=False, metadata={"help": "lambda for uniform loss"}
    )
    club_iters: int = field(
        default=5, metadata={"help": "lambda for uniform loss"}
    )
    club_weight: float = field(
        default=0.0001, metadata={"help": "lambda for uniform loss"}
    )
    club_target: bool = field(
        default=False, metadata={"help": "if regulation for target"}
    )
    club_warmup_iters: int = field(
        default=10, metadata={"help": "warm_up_iters for club"}
    )
    use_flash_attn: bool = False 
    train_vision_encoder: bool = False
    train_vision_projector: bool = False
    train_llm: bool = False
    train_parscale: bool = False
    vision_projector_lr: float = None 

    def __post_init__(self):
        super().__post_init__()
        if self.resume_from_checkpoint == "False":
            self.resume_from_checkpoint = False
        elif self.resume_from_checkpoint == "True":
            self.resume_from_checkpoint = True
        self.remove_unused_columns = False


@dataclass
class LoraArguments:
    use_lora: bool = True
    use_vision_lora: bool = True
    q_lora: bool = False
    lora_r: int = 64
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_weight_path: str = ""
    lora_bias: str = "none"
    vision_lora_r: int = 16
    vision_lora_alpha: int = 16

@dataclass
class FCArguments:
    r"""
    Arguments pertaining to which techniques we are going to fine-tuning with.
    """
    stage: Optional[Literal["classify_infer_crop", "classify_infer", "pt", "sft", "rm", "ppo", "retrieval", "reranker", "retrieval_infer_query", "retrieval_infer_item", "retrieval_hard_negative", "classify"]] = field(
        default="classify",
        metadata={"help": "Which stage will be performed in training."},
    )
    embed_dim: Optional[int] = field(
        default=512,
        metadata={"help": "embed dimension of output token"},
    )
    temperature: Optional[float] = field(
        default=0.02,
        metadata={"help": "temperature of the infonce loss"},
    )
    neg_gather_size: Optional[int] = field(
        default=1,
        metadata={
            "help": "how many gpus share negatives in the infonce loss. "
                    "1 = only in-device batch; 8 = share within a node; "
                    "-1 or world_size = share across all gpus. "
                    "candidate pool = 2 * neg_gather_size * per_device_train_batch_size"
        },
    )
    classify_num: Optional[int] = field(
        default=100000,
        metadata={"help": "Num of classification"},
    )
    classify_sample_rate: Optional[float] = field(
        default=0.05,
        metadata={"help": "Sample rate for partial fc"},
    )
    classify_use_knn: Optional[bool] = field(
        default=True,
        metadata={"help": "Use knn for partial fc"},
    )
    classify_global_knn: Optional[bool] = field(
        default=False,
        metadata={"help": "Use knn for partial fc"},
    )
    classify_knn_topk: Optional[int] = field(
        default=30,
        metadata={"help": "Knn topk for partial fc"},
    )
    classify_knn_chunk_size: Optional[int] = field(
        default=5000,
        metadata={"help": "knn chunk size for sim matrix calc: only open when global_knn"},
    )
    classify_knn_ratio: Optional[float] = field(
        default=0.5,
        metadata={"help": "Knn rate for partial fc"},
    )
    classify_knn_weight_cover_rate: Optional[float] = field(
        default=0.3,
        metadata={"help": "Knn rate for partial fc"},
    )
    classify_knn_update_after_save: Optional[bool] = field(
        default=True,
        metadata={"help": "knn_update_after_save"},
    )
    classify_loss_type: Optional[str] = field(
        default="arcface",
        metadata={"help": "loss function for partial fc"},
    )
    classify_ckpt_part: Optional[bool] = field(
        default=True,
        metadata={"help": "partial fc ckpt save load partly"},
    )
    classify_freeze_backbone: Optional[bool] = field(
        default=False,
        metadata={"help": "freeze backbone, only train partial fc"},
    )
    model_name_or_path_pfc: Optional[str] = field(
        default="/data/oss_bucket_0/guijiang/outputs/tongkuan/lamra_qwen2_5_fc_debug",
        metadata={"help": "pth to fc pth"},
    )
    model_name_or_path_proj: Optional[str] = field(
        default="/data/oss_bucket_0/guijiang/outputs/tongkuan/lamra_qwen2_5_fc_debug",
        metadata={"help": "pth to fc pth"},
    )
    proj_name: Optional[Literal["load_add_projector", "load_mlp_projector"]] = field(
        default="load_add_projector",
        metadata={"help": "Which stage will be performed in training."},
    )
    pfc_trainable: Optional[bool] = field(
        default=True,
        metadata={"help": "train pfc or not"},
    )
    def save_to_json(self, json_path: str):
        r"""Saves the content of this instance in JSON format inside `json_path`."""
        json_string = json.dumps(asdict(self), indent=2, sort_keys=True) + "\n"
        with open(json_path, "w", encoding="utf-8") as f:
            f.write(json_string)

    @classmethod
    def load_from_json(cls, json_path: str):
        r"""Creates an instance from the content of `json_path`."""
        with open(json_path, "r", encoding="utf-8") as f:
            text = f.read()

        return cls(**json.loads(text))