import json
from src.correlation import Variable, Correlation
import ast
import pandas as pd

def parse_json_block(text):
    # Find content between ```json and ```
    start = text.find('```json') + 7 
    end = text.find('```', start)
    
    if start == -1 or end == -1:
        raise ValueError("JSON block not found")

    json_str = text[start:end].strip()
    
    return json.loads(json_str)

def get_range(r_obs, error_bound, lower=-1, upper=1):
    return max(lower, r_obs - error_bound), min(upper, r_obs + error_bound)


def get_msg_for_a_correlation(data, pair_id, prompt, context=None):
    row = data.loc[data['pair_id'] == pair_id].iloc[0]
    original_row = row.to_dict()
    # print(original_row['var2_desc'])
    pair_id = original_row['pair_id']
    var1_data = {
        "attr": original_row['var1'],
        "table": original_row['dataset'],
        "table_desc": original_row['dataset description'],
        "var_desc": original_row['var1_desc']
    }
    var2_data = {
        "attr": original_row['var2'],
        "table": original_row['dataset'],
        "table_desc": original_row['dataset description'],
        "var_desc": original_row['var2_desc']
    }
    variable1 = Variable(**var1_data)
    variable2 = Variable(**var2_data)
    correlation = Correlation(variable1, variable2)
    msg = prompt.format(attr1=variable1.attr, attr2=variable2.attr, table=variable1.table, tbl_desc=variable1.table_desc, var1_desc=variable1.var_desc, var2_desc=variable2.var_desc, context=context)
    return msg, correlation, original_row

if __name__ == "__main__":
    text = """```json
    {
      "coefficient": 0.65,
      "standard deviation": 0.1
    }
    ```
    """
    print(parse_json_block(text))  


def add_understood_unclear_counts_from_str(df, json_col):
    """
    Parses stringified dicts using ast.literal_eval and adds per-row counts
    for 'Understood' and 'Unclear'.

    Parameters:
        df (pd.DataFrame): Your DataFrame.
        json_col (str): Column name with dict-like strings.

    Returns:
        pd.DataFrame: DataFrame with two new columns added.
    """
    def parse_and_count(val):
        try:
            d = ast.literal_eval(val) if isinstance(val, str) else val
        except Exception:
            d = {}
        understood = sum(1 for v in d.values() if v == "Understood")
        unclear = sum(1 for v in d.values() if v == "Unclear")
        return pd.Series([understood, unclear])

    df[["understood_count", "unclear_count"]] = df[json_col].apply(parse_and_count)
    return df