import os
import argparse
import json
import re
import jsonlines
from fractions import Fraction
from vllm import LLM, SamplingParams
import sys
import pandas as pd
MAX_INT = sys.maxsize
from huggingface_hub import login


def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        pass
    try:
        import unicodedata
        unicodedata.numeric(s)
        return True
    except (TypeError, ValueError):
        pass
    return False
import re
from fractions import Fraction

def extract_answer_number(completion):
    """
    Extract the first number (integer, decimal, fraction, or yes/no) from the model's response.
    """
    if "The answer is:" in completion:
        text = completion.split("The answer is:")[-1].strip().lower()
    else:
        text = completion.strip().lower()

    # Check for 'yes'/'no' response
    if "yes" in text:
        return "yes"
    elif "no" in text:
        return "no"

    # Extract only numbers (integer, decimal, fraction, all supported)
    match = re.search(r'[-+]?\d*\.?\d+(?:/\d+)?', text)
    
    if match:
        num_str = match.group()
        try:
            # Process fractions
            if '/' in num_str:
                frac = Fraction(num_str)
                # Use Fraction object for accurate calculation
                return float(frac)
            else:
                return float(num_str)
        except Exception:
            return None
    
    return None


def batch_data(data_list, batch_size=1):
    n = len(data_list) // batch_size
    batch_data = []
    for i in range(n-1):
        start = i * batch_size
        end = (i+1)*batch_size
        batch_data.append(data_list[start:end])

    last_start = (n-1) * batch_size
    last_end = MAX_INT
    batch_data.append(data_list[last_start:last_end])
    return batch_data


def finqa_test(model, data_path, start=0, end=MAX_INT, batch_size=1, tensor_parallel_size=1, seq_len=4096):
    INVALID_ANS = "[invalid]"
    qq=[]
    finqa_ins = []
    finqa_answers = []
    
    problem_prompt = (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Response: Let's think step by step."
    )
    print('promt =====', problem_prompt)
    with open(data_path,"r+", encoding="utf8") as f:
        for idx, item in enumerate(jsonlines.Reader(f)):
            qq.append(item["query"])
            temp_instr = problem_prompt.format(instruction=item["query"])
            finqa_ins.append(temp_instr)
            
            # Process response
            response = item['response']
            # Check if the response is a number or a string and process it appropriately
            if isinstance(response, (int, float)) or (isinstance(response, str) and response.replace('.', '', 1).replace('-', '', 1).isdigit()):
                # If the response is a number, convert it to float
                temp_ans = float(response)
            else:
                # If the response is a string (yes/no, etc.), keep it as is
                temp_ans = response
                
            finqa_answers.append(temp_ans)

    # finqa_ins = finqa_ins[start:end]
    # finqa_answers = finqa_answers[start:end]
    print('lenght ====', len(finqa_ins))
    batch_finqa_ins = batch_data(finqa_ins, batch_size=batch_size)

    stop_tokens = ["Question:", "Question", "USER:", "USER", "ASSISTANT:", "ASSISTANT", "Instruction:", "Instruction", "Response:", "Response"]
    sampling_params = SamplingParams(temperature=0 ,top_p=1.0, max_tokens=seq_len, stop=stop_tokens)
    print('sampleing =====', sampling_params)
    llm = LLM(model=model,tensor_parallel_size=tensor_parallel_size)

    print('#### Model was loaded !!!! ####')
    result = []
    res_completions = []
    output_lst =[]
    for idx, (prompt, prompt_answer) in enumerate(zip(batch_finqa_ins, finqa_answers)):
        if isinstance(prompt, list):
            pass
        else:
            prompt = [prompt]

        completions = llm.generate(prompt, sampling_params)
        for output in completions:
            prompt = output.prompt
            generated_text = output.outputs[0].text
            res_completions.append(generated_text)

    
    
    invalid_outputs = []
    gt=[]

    for idx, (prompt, completion, prompt_answer) in enumerate(zip(finqa_ins, res_completions, finqa_answers)):
        doc = {'question': prompt}
        y_pred = extract_answer_number(completion)
        gt.append(prompt_answer)
        if y_pred != None:
            # If both are strings (yes/no)
            if isinstance(y_pred, str) and isinstance(prompt_answer, str):
                is_correct = y_pred.lower() == prompt_answer.lower()
            # If both are numbers
            elif isinstance(y_pred, (int, float)) or isinstance(prompt_answer, (int, float)):
                try:
                    # Round for comparison (4 decimal places)
                    decimal_places = 4
                    rounded_pred = round(float(y_pred), decimal_places)
                    rounded_answer = round(float(prompt_answer), decimal_places)
                    is_correct = rounded_pred == rounded_answer
                except (ValueError, TypeError):
                    is_correct = False
            else:
                is_correct = False
                
            result.append(is_correct)
        else:
            result.append(False)
            temp = {'question': prompt, 'output': completion, 'answer': prompt_answer}
            invalid_outputs.append(temp)
    acc = sum(result) / len(result)
    
    df=pd.DataFrame({'query':qq ,'answer':res_completions, 'GT': gt})

    # Save CSV file in absolute path
    output_path = f'./{args.outdir}.csv'
    df.to_csv(output_path, index=False)
    print(f'Results saved to: {output_path}')
    
        
    print('='*50)
    print('Model:', args.outdir)
    print('Generation was finished')
    print('='*50)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str)  # model path
    parser.add_argument("--outdir", type=str,default='output_answer')  # model path
    parser.add_argument("--data_file", type=str, default='')  # data path
    parser.add_argument("--start", type=int, default=0) #start index
    parser.add_argument("--end", type=int, default=MAX_INT)  # end index
    parser.add_argument("--seq_len", type=int, default=4096)
    parser.add_argument("--batch_size", type=int, default=400)  # batch_size
    parser.add_argument("--tensor_parallel_size", type=int, default=4)  # tensor_parallel_size
    return parser.parse_args()
if __name__ == "__main__":
    args = parse_args()
    finqa_test(model=args.model, data_path=args.data_file, start=args.start, end=args.end, batch_size=args.batch_size, tensor_parallel_size=args.tensor_parallel_size, seq_len=args.seq_len)
