from modules import *
import os, sys
import numpy as np
from tqdm import tqdm
import torch
from torch import nn
from config import CFG
import utils
import json
import pandas as pd
import pickle
from rdkit import Chem
from rdkit.Chem import inchi

def smiles_to_inchikey(smiles, nostereo=True):
    try:
        # 将SMILES转换为分子对象
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        
        if nostereo:
            options = "-SNon"
            inchi_string = inchi.MolToInchi(mol, options=options)
        else:
            inchi_string = inchi.MolToInchi(mol)

        if not inchi_string:
            return None
        
        inchikey = inchi.InchiToInchiKey(inchi_string)

        return inchikey
        
    except Exception as e:
        print(f"转换失败: {e}")
        return None

def calc_mol_embeddings(model, smis, cfg):
    model.eval()
    fp_featsl = []
    gnn_featsl = []
    fm_featsl = []
    valid_smis = []

    for smil in smis:
        smi = smil[1]
        try:
            if 'gnn' in cfg.mol_encoder:
                gnn_feats = utils.mol_graph_featurizer(smi)
                gnn_featsl.append(gnn_feats)
            if 'fp' in cfg.mol_encoder:
                fp_feats = utils.mol_fp_encoder(smi, tp=cfg.fptype, nbits=cfg.mol_embedding_dim).to(cfg.device)
                fp_featsl.append(fp_feats)
            if 'fm' in cfg.mol_encoder:
                fm_feats = utils.smi2fmvec(smi).to(cfg.device)
                fm_featsl.append(fm_feats)
            valid_smis.append(smil)
        except Exception as e:
            print(smi, e)
            continue
    
    mol_feat_list = []
    if 'gnn' in cfg.mol_encoder:
        vl, al, msl = [], [], []
        bat = {}
        for b in gnn_featsl:
            if 'V' in b:
                vl.append(b['V'])
            if 'A' in b:
                al.append(b['A'])
            if 'mol_size' in b:
                msl.append(b['mol_size'])
        
        vl1, al1 = [], []
        if vl and al and msl:
            max_n = max(map(lambda x:x.shape[0], vl))
            for v in vl:
                vl1.append(utils.pad_V(v, max_n))
            for a in al:
                al1.append(utils.pad_A(a, max_n))

        bat['V'] = torch.stack(vl1).to(cfg.device)
        bat['A'] = torch.stack(al1).to(cfg.device)
        bat['mol_size'] = torch.cat(msl, dim=0).to(cfg.device)

        mol_feat_list.append(model.mol_gnn_encoder(bat))

        del bat

    if 'fp' in cfg.mol_encoder:
        mol_feat_list.append(torch.stack(fp_featsl).to(cfg.device))

    if 'fm' in cfg.mol_encoder:
        mol_feat_list.append(torch.stack(fm_featsl).to(cfg.device))
        
    if len(mol_feat_list) > 1:
        mol_features = torch.cat(mol_feat_list, dim=1).to(cfg.device)
    else:
        mol_features = mol_feat_list[0].to(cfg.device)

    with torch.no_grad():
        mol_embeddings = model.mol_projection(mol_features)

    del mol_features, mol_feat_list

    return mol_embeddings, valid_smis

def find_matches(model, ms, smis, cfg, n=10, batch_size=64):
    model.eval()
    with torch.no_grad():
        ms_features = utils.ms_binner(ms, min_mz=cfg.min_mz, max_mz=cfg.max_mz, bin_size=cfg.bin_size, add_nl=cfg.add_nl, binary_intn=cfg.binary_intn).to(cfg.device)
        ms_features = ms_features.unsqueeze(0)
        ms_embeddings = model.ms_projection(ms_features)
        ms_embeddings_n = F.normalize(ms_embeddings, p=2, dim=1)
    
    # 分批计算相似度并维护top-k
    all_similarities = []
    all_valid_smis = []
    
    # 收集所有分子embedding
    all_embeddings = []
    for i in tqdm(range(0, len(smis), batch_size)):
        batch_smis = smis[i:i+batch_size]
        batch_embeddings, valid_smis = calc_mol_embeddings(model, batch_smis, cfg)
        all_embeddings.append(batch_embeddings)
        all_valid_smis.extend(valid_smis)

        del batch_embeddings
    
    # 全局归一化
    all_embeddings = torch.cat(all_embeddings, dim=0)
    all_embeddings_n = F.normalize(all_embeddings, p=2, dim=1)
    
    # 计算相似度
    similarities = F.cosine_similarity(all_embeddings_n, ms_embeddings_n, dim=1)
    #print('all_embeddings_n.shape', all_embeddings_n.shape, ms_embeddings.shape, len(all_valid_smis), similarities.shape)
    
    if n == -1 or n > len(all_valid_smis):
        n = len(all_valid_smis)
    
    values, topk_indices = torch.topk(similarities, n)
    
    topk_indices_list = topk_indices.cpu().tolist()
    #print(len(topk_indices_list), len(all_valid_smis), len(similarities))
    matchsmis = [all_valid_smis[idx] for idx in topk_indices_list]

    return matchsmis, values.cpu().numpy()*100, topk_indices_list

def calc(models, datal, cfg):
    dicall = {}
    coridxd = {}

    for idx, model in enumerate(models):
        for nn, data in enumerate(datal):
            print(f'Calculating {nn}-th MS...')

            try:
                smis, scores, indices = find_matches(model, data['ms'], data['candidates'], cfg, 50)
            except Exception as e:
                print(131, e)
                continue

            dic = {}
            for n, smil in enumerate(smis):
                smi = smil[1]
                if smi in dic:
                    dic[smi]['score'] = scores[n]
                    dic[smi]['iscor'] = smis[n][-1]
                    dic[smi]['idx'] = smis[n][0]
                else:
                    dic[smi] = {'score': scores[n], 'iscor': smis[n][-1], 'idx': smis[n][0]}

            # 计算去除立体构型分子的inchikey，由于质谱很难区分立体构型，我们认为分子的不同立体构型都算正确匹配
            ikey = smiles_to_inchikey(data['smiles'], True)
            if ikey is None:
                ikey = data['ikey']

            if ikey in dicall:
                for k, v in dic.items():
                    if k in dicall[ikey]:
                        dicall[ikey][k]['score'] += v['score']
                        dicall[ikey][k]['score'] /= 2
                    else:
                        dicall[ikey][k] = v
            else:
                dicall[ikey] = dic

    for ikey, dic in dicall.items():
        smis   = [k for k in dic.keys()]
        scorel = [d['score'] for d in dic.values()]
        iscorl = [d['iscor'] for d in dic.values()]
        indexl = [d['idx'] for d in dic.values()]

        scoretsor = torch.tensor(scorel)
        n = 100
        if n > len(scorel):
            n = len(scorel)

        values, indices = torch.topk(scoretsor, n)
        
        # 修复：将张量转换为Python列表
        indices_list = indices.cpu().tolist()
        
        scorel = values.cpu().numpy()
        smis   = [smis[i] for i in indices_list]
        iscorl = [iscorl[i] for i in indices_list]
        indexl = [indexl[i] for i in indices_list]

        try:
            i = iscorl.index(True)
            k = 'Hit %.3d' %(i+1)
            if k in coridxd:
                coridxd[k] += 1
            else:
                coridxd[k] = 1
        except:
            pass

    ks = sorted(list(coridxd.keys()))
    dc = {}
    sumtop3 = 0

    for k in ks:
        dc[k] = [coridxd[k]]
        if k in ['Hit 001', 'Hit 002', 'Hit 003']:
            sumtop3 += coridxd[k]

    for i in range(100):
        k = 'Hit %.3d' %(i+1)
        if not k in dc:
            dc[k] = [0]

    return sumtop3, dc, dicall

def calc_rank(dicall):
    rankd = {}

    for ikey, dic in dicall.items():
        smis   = [k for k in dic.keys()]
        scorel = [d['score'] for d in dic.values()]
        iscorl = [d['iscor'] for d in dic.values()]
        indexl = [d['idx'] for d in dic.values()]

        scoretsor = torch.tensor(scorel)
        n = 100
        if n > len(scorel):
            n = len(scorel)

        values, indices = torch.topk(scoretsor, n)

        scorel = values
        smis   = [smis[i] for i in indices]
        iscorl = [iscorl[i] for i in indices]
        indexl = [indexl[i] for i in indices]

        sl = []
        for n, smi in enumerate(smis):
            sl.append(f'{scorel[n]}:{smi}:{smiles_to_inchikey(smi)}')

        try:
            i = iscorl.index(True)
            rankd[ikey] = {'Hit': i+1, 'Rank': sl}
        except:
            pass

    return rankd

def predict(modelfnl, datal, datafn=''):
    maxtop3 = 0
    maxoutt = ''

    for fn in modelfnl:
        d = torch.load(fn)
        CFG.load(d['config'])
        print(d['config'])
        CFG.save('', True)

        model = FragSimiModel(CFG).to(CFG.device)
        model.load_state_dict(d['state_dict'])

        sumtop3, dc, dicall = calc([model], datal, CFG)

        sumtop10 = 0
        for k in ['Hit %.3d' %(i+1) for i in range(10)]:
            if k in dc:
                sumtop10 += dc[k][0]

        sumtop50 = 0
        for k in ['Hit %.3d' %(i+1) for i in range(50)]:
            if k in dc:
                sumtop50 += dc[k][0]

        tops = {}
        for i in range(100):
            k = 'Hit %.3d' %(i+1)
            key = k.replace('Hit', 'Top')
            if not key in tops:
                tops[key] = [0]
            if k in dc:
                for n in range(i+1):
                    kk = 'Hit %.3d' %(n+1)
                    if kk in dc:
                        tops[key][0] += dc[kk][0]

        outt = f'Top1: {dc.setdefault("Hit 001", [0])[0]}, top3: {sumtop3}, top10: {sumtop10}, top50: {sumtop50} of {len(datal)}'

        if sumtop3 > maxtop3:
            maxtop3 = sumtop3
            maxoutt = outt

        basefn = fn.replace('.pth', f'-{os.path.basename(datafn).split(".")[0]}')
        rank = calc_rank(dicall)
        json.dump(rank, open(basefn + '-predict-rank.json', 'w'), indent=2)

        df = pd.DataFrame(tops)
        df.to_csv(basefn + '-predict-summary.csv', index=False)

    return maxoutt, maxtop3

def main(datafn, fnl):
    outl = []

    datal = json.load(open(datafn))

    n = 0
    for n, fn in enumerate(fnl):
        out, _ = predict([fn], datal, datafn)
        print(out, os.path.basename(fn))
        outl.append(out)

    print(outl)

if __name__ == '__main__':
    import time
    t0 = time.time()
    main(sys.argv[1], sys.argv[2:])
    print(300, time.time()-t0)