from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from prompt import *
from utils import call, permutate

import os
import json
import pandas as pd
import re

load_dotenv()

class Validation():

    def __init__(self, llm):

        if llm.startswith('gpt'):
            self.llm = ChatOpenAI(temperature=0, model_name=llm)
        elif llm.startswith('gemini'):
            self.llm = ChatGoogleGenerativeAI(temperature=0, model=llm)
        else:
            self.llm = ChatOpenAI(temperature=0, model_name=llm, api_key=os.environ['PERPLEXITY_API_KEY'], base_url="https://api.perplexity.ai")

    def validate(self, df, api):

        df = df.fillna('')
        df['Genes'] = df['Genes'].str.replace(' ', '').str.upper()
        df['SNPs'] = df['SNPs'].str.lower()

        # Check if there is two gene names
        sym = [',', '/', '|', '-']
        for i in df.index:
            gene = df.loc[i, 'Genes']
            for s in sym:
                if s in gene:
                    genes = gene.split(s)
                    df.loc[i + 0.1], df.loc[i + 0.9] = df.loc[i], df.loc[i]
                    df = df.sort_index().reset_index(drop=True)
                    df.loc[i + 1, 'Genes'], df.loc[i + 2, 'Genes'] = genes[0], s.join(genes[1:])
                    break

        # Check if there is SNPs without 'rs'
        for i in df.index:
            safe = True
            snp = df.loc[i, 'SNPs']
            snp = snp.replace('l', '1')
            if re.fullmatch('rs(\d)+|', snp):
                pass
            elif re.fullmatch('ts(\d)+', snp):
                snp = 'r' + snp[1:]
            elif re.fullmatch('s(\d)+', snp):
                snp = 'r' + snp
            elif re.fullmatch('(\d)+', snp):
                snp = 'rs' + snp
            else:
                safe = False
                df = df.drop(i)

            if safe:
                df.loc[i, 'SNPs'] = snp

        df.reset_index(drop=True, inplace=True)
        df_clean = df.copy()

        # Validate genes and SNPs with APIs
        if api:
            dbsnp = {}
            for i in df.index:
                snp = df.loc[i, 'SNPs']
                gene = df.loc[i, 'Genes']

                if snp not in dbsnp:
                    res = call(f'https://www.ebi.ac.uk/gwas/rest/api/singleNucleotidePolymorphisms/{snp}/')
                    try:
                        res = res.json()
                        dbsnp[snp] = [r['gene']['geneName'] for r in res['genomicContexts']]
                    except:
                        print("Error at first API", e)
                        dbsnp[snp] = []

                    try:
                        res = call(f'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=snp&retmode=json&id={snp[2:]}').json()['result'][snp[2:]]
                        if 'error' not in res:
                            dbsnp[snp].extend([r['name'] for r in res['genes']])
                    except Exception as e:
                        print("Error at second API", e)
                        pass

                    dbsnp[snp] = list(set(dbsnp[snp]))

                if gene not in dbsnp[snp]:
                    for other in permutate(gene):
                        if other in dbsnp[snp]:
                            df.loc[i, 'Genes'] = other
                            print(f'{gene} corrected to {other}')
                            break
                    else:
                        df = df.drop(i)

        df.reset_index(drop=True, inplace=True)
        df_no_llm = df.copy()

        # Validate genes and diseases with LLM (for each 50 rows)
        idx = 0
        results = []

        while True:
            json_table = df[['Genes', 'SNPs', 'Diseases']][idx:idx+50].to_json(orient='records')
            str_json_table = json.dumps(json.loads(json_table), indent=2)

            result = self.llm.invoke(input=prompt_validation.format(str_json_table)).content
            print('val', idx)
            print(result)

            result = result[result.find('['):result.rfind(']')+1]
            try:
                result = eval(result)
            except SyntaxError:
                result = []

            results.extend(result)
            idx += 50
            if idx not in df.index:
                break

        df = pd.DataFrame(results)
        df = df.merge(df_no_llm.head(1).drop(['Genes', 'SNPs', 'Diseases'], axis=1), 'cross')

        return df, df_no_llm, df_clean