#!/usr/bin/env python3
"""python3 score.py raw/<file>.jsonl : brand mention / first-mention / own-site citation counts -> results.csv"""
import csv, json, re, sys, urllib.parse, collections
# brand: (regex over answer text, own domains)
B = {
 'Nanchang Mineral Systems (NMS)': (r'Nanchang Mineral|\bNMS\b|南昌矿机', ['nmsindustries.com']),
 'Liming Heavy Industry': (r'Liming|黎明重工', ['limingco.com', 'liming-global.com']),
 'SANME': (r'SANME|山美', ['sanmecrusher.com', 'sanmecrushers.com', 'sanmecrushing.com', 'sanmemachine.com']),
 'SBM (Shibang)': (r'\bSBM\b|Shibang|世邦', ['sbmchina.com', 'sbmvietnam.com', 'sbmcn.com', 'shibangchina.com']),
 'ZENITH': (r'ZENITH|\bZenith\b', ['zenithcrusher.com', 'zenith-mills.com', 'zenithmining.com']),
 'JXSC': (r'JXSC', ['jxscmachine.com', 'jxscmineral.com', 'jxscmining.com', 'jxscteam.com']),
 'CITIC Heavy Industries (incl. CITIC-TLC)': (r'CITIC|中信重工|CIC Luoyang', ['citicheavyindustries.com', 'citictlc.com', 'en.cichmc.com']),
 'Xinhai Mining': (r'Xinhai|鑫海', ['xinhaimining.com']),
 'XCMG': (r'XCMG|徐工', ['xcmg.co', 'xcmg.net.au', 'xcmg.com']),
 'SANY': (r'\bSANY\b|\bSany\b|三一', ['sanyglobal.com']),
 'Zoomlion': (r'Zoomlion|中联重科', ['zoomlion.com']),
 'NFLG (South Highway Machinery)': (r'NFLG|South Highway Machinery|南方路机', ['nflg.com', 'nflg.ae']),
 'Northern Heavy Industries (NHI / SNM)': (r'Northern Heavy|\bNHI\b|\bSNM\b|北方重工', ['nhi-international.com', 'nhi.com.cn', 'snmhi.com']),
 'Fote (FTM)': (r'Fote|\bFTM\b', ['ftmmachinery.com']),
 'DSMAC (Dingsheng)': (r'DSMAC|Dingsheng', ['dscrusher.com']),
 'Prominer': (r'Prominer', ['prominetech.com']),
 'BGRIMM': (r'BGRIMM', ['english.bgrimm.com']),
 'China ENFI': (r'\bENFI\b', []),
 'Sinosteel': (r'Sinosteel', ['entech.sinosteel.com']),
 'Anzhong (鞍重)': (r'Anzhong|鞍重', ['anzhongtechnology.com']),
 'Kminda': (r'Kminda', ['kminda.com']),
 'MESDA': (r'MESDA', ['mesda.com', 'mesdagroup.com']),
 'Shanbao': (r'SHANBAO', ['shanbaomachinery.com']),
 'Guilin Hongcheng': (r'Hongcheng|鸿程', ['new.hcmilling.com']),
 'Shanghai Clirik': (r'Clirik', ['clirik.com']),
 'Zhengzhou YIFAN': (r'YIFAN', ['yfcrusher.com']),
 'Shanghai Kinglink': (r'Kinglink', []),
 'Zhejiang Shuangjin': (r'Shuangjin|双金', []),
 'Baichy': (r'Baichy', ['baichychina.com']),
 'Hongji': (r'Hongji', ['chinahongji.com']),
 'Luoyang Dahua': (r'Luoyang Dahua|\bDahua\b', ['lydhcrusher.com']),
 'Dasen Mining': (r'\bDasen\b', ['dasenmining.com']),
 'LiuGong': (r'LiuGong|柳工', ['liugong.com']),
 'AIMIX': (r'AIMIX', ['aimixgroup.com']),
 'Yunnan Kreat': (r'\bKreat\b', ['kreatgroup.com']),
 'Henan ZK Corp': (r'ZK Corp|Henan ZK|ZK Ecotech', ['zkcomp.com']),
 'Yantai Jinpeng': (r'Jinpeng', ['jinpengmining.com']),
 'Beijing HOT Mining': (r'HOT Mining', ['hotminingepc.com']),
 'Aury (奥瑞)': (r'\bAury\b|奥瑞', ['aury-china.com']),
 'DASWELL': (r'DASWELL', []),
 'Camelway': (r'Camelway', ['camelway.com']),
 'Henan Hongxing (红星)': (r'Hongxing|红星机器', ['hxjq.com', 'hxjqchina.com', 'hxjq-machine.com']),
 'ZMJ (Zhengzhou Coal Mining Machinery)': (r'\bZMJ\b|Zhengzhou Coal Mining', ['zmj.com']),
 'Taiyuan Heavy Industry (TYHI)': (r'Taiyuan Heavy|\bTYHI\b', []),
 'Sinoma': (r'Sinoma', ['sinomaly-china.com']),
 'Huate Magnet': (r'Huate', ['huatemagnets.com']),
 'Hengchang Mining Machinery': (r'Hengchang', ['hengchangmachinery.com']),
}
L = [json.loads(l) for l in open(sys.argv[1])]
n = len(L); rows = collections.defaultdict(lambda: {'mention': 0, 'first': 0, 'cited': 0, 'ids': []})
for d in L:
    a = d['answer']; first = None
    doms = {urllib.parse.urlsplit(u).netloc.replace('www.', '') for u in d['links'] or []}
    for b, (rx, ds) in B.items():
        m = re.search(rx, a, re.I if b not in ('SANY', 'ZENITH') else 0)
        if m:
            rows[b]['mention'] += 1; rows[b]['ids'].append(d['id'])
            if first is None or m.start() < first[1]: first = (b, m.start())
        if any(h == x or h.endswith('.' + x) for h in doms for x in ds): rows[b]['cited'] += 1  # subdomains count
    if first: rows[first[0]]['first'] += 1
out = sorted(rows.items(), key=lambda x: (-x[1]['mention'], -x[1]['first']))
with open('results.csv', 'w', newline='') as f:
    w = csv.writer(f); w.writerow(['rank', 'brand', 'mentioned_in', 'mention_rate', 'first_mentioned_in', 'first_rate', 'own_site_cited_in', 'prompt_ids'])
    for i, (b, r) in enumerate(out, 1):
        w.writerow([i, b, r['mention'], f"{r['mention']/n:.0%}", r['first'], f"{r['first']/n:.0%}", r['cited'], ' '.join(r['ids'])])
        print(f"{i:>2} {b:42} {r['mention']:>2}/{n} {r['mention']/n:>4.0%}  first {r['first']:>2}  cited {r['cited']:>2}")