#!/usr/bin/env python3
"""python3 score.py raw/<file>.jsonl [more.jsonl ...] : brand mention / first-mention / own-site citation counts -> results.csv
Only brands verified as Chinese companies are scored. Seen in answers but not verified as Chinese (left out):
Himiway, Velotric, Heybike, Aventon, Lectric, DTTZH, TYCORUN, TTFAR, BENLG, Totem, Super Soco (Vmoto). Foreign component brands (Bosch, Shimano, Brose) are not scored."""
import csv, json, re, sys, urllib.parse, collections
CS = {'ADO', 'EVE Energy', 'DYU', 'HIMO', 'JSL'}  # case-sensitive (common words otherwise)
# brand: (regex over answer text, own domains)
B = {
 # two-wheeler makers
 'Yadea': (r'\bYadea', ['yadea.com', 'yadea.com.cn']),
 'AIMA': (r'\bAima\b', ['aimatech.com', 'aimaebike.com']),
 'TAILG': (r'\bTailg\b', ['tailg.com', 'tailg.com.cn']),
 'Sunra (Xinri)': (r'\bSunra\b|\bXinri\b', ['sunraev.com', 'sunra.com']),
 'Luyuan': (r'\bLuyuan\b', ['luyuangroup.com', 'luyuanvehicle.com', 'luyuan.cn']),
 'NIU': (r'\bNIU\b|\bNiu Technologies', ['niu.com']),
 'Segway-Ninebot': (r'\bSegway\b|\bNinebot\b', ['segway.com', 'ninebot.com']),
 'Xiaomi': (r'\bXiaomi\b', ['mi.com']),
 'HIMO': (r'\bHIMO\b|\bHimo\b', ['himo.com']),
 'Zongshen': (r'\bZongshen\b', ['zongshen.com', 'zongshenmotor.com']),
 'Golden Wheel Group': (r'Golden Wheel', ['goldenwheelgroup.com']),
 'Sur-Ron': (r'Sur-?Ron', ['sur-ron.com', 'sur-ronusa.com']),
 'Talaria': (r'\bTalaria\b', ['talaria.bike', 'talariabike.com']),
 'ZEEHO': (r'\bZEEHO\b', ['zeehoev.com']),
 # e-bike brands
 'Fiido': (r'\bFiido\b', ['fiido.com']),
 'ENGWE': (r'\bENGWE\b', ['engwe.com', 'engwe-bikes.com', 'engwe-bikes-eu.com']),
 'DYU': (r'\bDYU\b', ['dyucycle.com']),
 'ADO': (r'\bADO\b', ['adoebike.com', 'adoebike.jp']),
 'Tenways': (r'\bTenways\b', ['tenways.com']),
 'Urtopia': (r'\bUrtopia\b', ['urtopia.com']),
 'DJI Avinox / Amflow': (r'\bAvinox\b|\bAmflow\b|\bDJI\b', ['amflowbikes.com', 'dji.com']),
 'LightCarbon': (r'LightCarbon', ['lightcarbon.com']),
 'JOBO': (r'\bJOBO\b|\bJobo\b', ['jobobike.com', 'jobobikes.com', 'joboev.com']),
 'SAMEBIKE': (r'SAMEBIKE|Samebike', ['samebike.com', 'gzsamebike.com']),
 'Leebike': (r'\bLeebike\b', ['leebike.com']),
 'Haidong': (r'Haidong', ['hd-ebike.com']),
 'VeloMove': (r'VeloMove|Velomove', ['velomove.com']),
 'Cemoto': (r'\bCemoto\b', ['cemotobike.com']),
 'Coswheel': (r'\bCoswheel\b', ['coswheel.com']),
 'United Mobility (Wuxi)': (r'United Mobility', ['unitedebike.com']),
 'JSL (Zhejiang)': (r'\bJSL\b', ['jslebikes.com']),
 # drive systems and batteries
 'Bafang': (r'\bBafang\b', ['bafang-e.com']),
 'Ananda': (r'\bAnanda\b', ['ananda-drive.com']),
 'Shengyi': (r'\bShengyi\b', ['szshengyi.com']),
 'Tongsheng': (r'\bTongsheng\b', ['tongsheng-e.com']),
 'MXUS': (r'\bMXUS\b', ['mxus.com', 'mxusmotor.com']),
 'Dapu': (r'\bDapu\b', ['dapumotor.com']),
 'Phylion': (r'\bPhylion\b', ['phylion.com']),
 'CATL': (r'\bCATL\b', ['catl.com']),
 'EVE Energy': (r'\bEVE\b', ['evebattery.com']),
}
L = [json.loads(l) for f in sys.argv[1:] for l in open(f)]
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, 0 if b in CS else re.I)
        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:24} {r['mention']:>3}/{n} {r['mention']/n:>4.0%}  first {r['first']:>2}  cited {r['cited']:>2}")