#!/usr/bin/env python3
"""Shift every date in a patient's JSONL files by one offset (a whole number of years, a multiple of 4, plus optional
days); intervals between dates are preserved."""
import argparse
import datetime as dt
import json
import random
import re
from collections import Counter
from pathlib import Path

MONTHS = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October",
          "November", "December"]
MON = r"(?P<mon{n}>Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|June?|July?|Aug(?:ust)?|Sep(?:t(?:ember)?)?" \
      r"|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)"
Y = r"(?:19|20)\d\d"
UNIT = (r"(?!\s?(?!U/A\b|L-[a-z])"
        r"(?:mg|mcg|g|kg|ml|mL|cc|L|IU|U|units?|kcal|cal|mmHg|mm|cm|%|tabs?|tablets?|caps?|times|x)\b)")
END = r"(?!\d|[/.]\d)"
RX = re.compile("|".join([
    rf"(?P<ymd>(?<![\d/])(?<!\d\.)(?P<y1>{Y})(?P<s1>[-/.])(?P<m1>\d{{1,2}})(?P=s1)(?P<d1>\d{{1,2}}){END})",
    rf"(?P<mdyn>(?<![\d/])(?<!\d\.)(?P<m12>0?[1-9]|1[0-2])/(?P<d12>0?[1-9]|[12]\d|3[01])/(?P<y12>{Y})(?!\d|[/.]\d))",
    rf"(?P<my2>(?<![\d/])(?<!\d\.)(?P<m13>0?[1-9]|1[0-2])/(?P<y13>{Y})(?!\d|[/.]\d))",
    rf"(?P<c8>(?<!\d)(?P<y2>{Y})(?P<m2>0[1-9]|1[0-2])(?P<d2>0[1-9]|[12]\d|3[01])(?!\d))",
    rf"(?P<mdy>{MON.format(n=1)}\.? (?P<d3>\d{{1,2}})(?P<s3>,? )(?P<y3>{Y})(?!\d))",
    rf"(?P<dmy>(?<!\d)(?P<d4>\d{{1,2}}) {MON.format(n=2)}(?P<s4>\.?,? )(?P<y4>{Y})(?!\d))",
    rf"(?P<yr>(?<![A-Za-z\d/])(?<!\d\.)(?P<y9>{Y})(?P<s9>-|/|\u2013)(?P<y10>{Y})(?![A-Za-z\d]|[/.]\d){UNIT})",
    rf"(?P<ym>(?<![\d/])(?<!\d\.)(?P<y5>{Y})(?P<s5>[-/.])(?P<m5>0?[1-9]|1[0-2]){END}(?!-\d{{1,2}}(?!\d)))",
    rf"(?P<my>{MON.format(n=3)}(?P<s6>\.?,? )(?P<y6>{Y})(?!\d))",
    rf"(?P<md>{MON.format(n=4)}\.? (?P<d7>\d{{1,2}})(?:st|nd|rd|th)?(?![\d,]))",
    rf"(?P<yx>(?<![A-Za-z\d/$])(?<!\d\.)(?P<y11>{Y})(?=[-/.]\d))",
    rf"(?P<y>(?<![A-Za-z\d/$])(?<!\d\.)(?P<y8>19[4-9]\d|20[0-2]\d)"
    rf"(?!(?!year)[A-Za-z]|\d|[/.]\d|-\d{{1,2}}(?!\d)){UNIT})",
]))
KINDS = ("ymd", "mdyn", "my2", "c8", "mdy", "dmy", "yr", "ym", "my", "md", "yx", "y")
MD_NUM = re.compile(r"(?<![\d/.])(?:0?[1-9]|1[0-2])/(?:0?[1-9]|[12]\d|3[01])(?![\d/])")


class ShiftError(Exception):
    """A real date that the offset cannot move by the same number of days as every other date."""


def month_no(name):
    return [m[:3] for m in MONTHS].index(name[:3]) + 1


class Shift:
    def __init__(self, years, days):
        if years % 4:
            raise ValueError(f"year offset {years} is not a multiple of 4")
        self.years, self.days = years, days
        self.delta = None
        self.n = Counter()

    def date(self, y, m, d):
        """Shift one date. Text that is not a real date raises ValueError and is left as it is; a real date that
        would move by a different number of days than the others (29 February into a year without it, or an
        offset across 1900 or 2100) raises ShiftError and stops the run."""
        src = dt.date(y, m, d)
        try:
            x = dt.date(y + self.years, m, d) + dt.timedelta(days=self.days)
        except ValueError:
            raise ShiftError(f"{src.isoformat()} has no counterpart {self.years:+d} years away") from None
        if self.delta is None:
            self.delta = (x - src).days
        elif (x - src).days != self.delta:
            raise ShiftError(f"{src.isoformat()} would move by {(x - src).days} days instead of {self.delta}")
        return x.year, x.month, x.day

    def __call__(self, s):
        return RX.sub(self.one, s)

    def one(self, m):
        k = next(x for x in KINDS if m.group(x))
        g = m.group
        try:
            if k == "ymd":
                y, mo, d = self.date(int(g("y1")), int(g("m1")), int(g("d1")))
                pad_m, pad_d = len(g("m1")), len(g("d1"))
                out = f"{y}{g('s1')}{mo:0{pad_m}d}{g('s1')}{d:0{pad_d}d}"
            elif k == "mdyn":
                y, mo, d = self.date(int(g("y12")), int(g("m12")), int(g("d12")))
                out = f"{mo:0{len(g('m12'))}d}/{d:0{len(g('d12'))}d}/{y}"
            elif k == "my2":
                out = f"{g('m13')}/{int(g('y13')) + self.years}"
            elif k == "c8":
                y, mo, d = self.date(int(g("y2")), int(g("m2")), int(g("d2")))
                out = f"{y}{mo:02d}{d:02d}"
            elif k == "mdy":
                y, mo, d = self.date(int(g("y3")), month_no(g("mon1")), int(g("d3")))
                out = (f"{self.mname(g('mon1'), mo)}{g(0)[len(g('mon1')):g(0).index(' ')]} "
                       f"{d:0{len(g('d3'))}d}{g('s3')}{y}")
            elif k == "dmy":
                y, mo, d = self.date(int(g("y4")), month_no(g("mon2")), int(g("d4")))
                out = f"{d:0{len(g('d4'))}d} {self.mname(g('mon2'), mo)}{g('s4')}{y}"
            elif k == "yx":
                out = str(int(g("y11")) + self.years)
            elif k == "yr":
                out = f"{int(g('y9')) + self.years}{g('s9')}{int(g('y10')) + self.years}"
            elif k == "ym":
                out = f"{int(g('y5')) + self.years}{g('s5')}{g('m5')}"
            elif k == "my":
                out = f"{g('mon3')}{g('s6')}{int(g('y6')) + self.years}"
            elif k == "md":
                out = g(0)
            else:
                out = str(int(g("y8")) + self.years)
        except ValueError:
            self.n["invalid " + k] += 1
            return g(0)
        self.n[k] += 1
        return out

    @staticmethod
    def mname(orig, mo):
        if mo == month_no(orig):
            return orig
        return MONTHS[mo - 1] if orig in MONTHS else MONTHS[mo - 1][:3]


def walk(o, f):
    if isinstance(o, dict):
        return {k: walk(v, f) for k, v in o.items()}
    if isinstance(o, list):
        return [walk(v, f) for v in o]
    return f(o) if isinstance(o, str) else o


def year_offsets(s):
    """Parse LO,HI into the multiples of 4 from LO to HI; any other offset is refused."""
    lo, hi = map(int, s.split(","))
    if lo > hi or lo % 4 or hi % 4:
        raise argparse.ArgumentTypeError(f"{s}: LO and HI must be multiples of 4 with LO <= HI")
    return list(range(lo, hi + 1, 4))


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--src", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--offsets", required=True, type=Path, help="TSV file that receives the offset of each patient")
    ap.add_argument("--years", type=year_offsets, required=True,
                    help="--years=LO,HI: each patient's year offset is a multiple of 4 from LO to HI, so leap days "
                         "and day intervals are kept; LO and HI must be multiples of 4")
    ap.add_argument("--max-days", type=int, default=0,
                    help="N (default 0): also add one day offset from -N to N to every full date of a patient; "
                         "year-month, year-only and month/day dates without a year are not shifted by days")
    ap.add_argument("--seed", type=int, required=True)
    ap.add_argument("--patients", nargs="*")
    a = ap.parse_args()
    if a.out.resolve() == a.src.resolve():
        raise SystemExit("--out must differ from --src")
    pids = a.patients or sorted(p.name for p in a.src.iterdir() if p.is_dir())
    total, rows = Counter(), []
    for pid in pids:
        rng = random.Random(f"{a.seed}:{pid}")
        sh = Shift(rng.choice(a.years), rng.randint(-a.max_days, a.max_days))
        (a.out / pid).mkdir(parents=True, exist_ok=True)
        no_year = 0
        for src in sorted((a.src / pid).glob("*.jsonl")):
            with open(src, encoding="utf-8") as fi, open(a.out / pid / src.name, "w", encoding="utf-8") as fo:
                for line in fi:
                    r = json.loads(line)
                    if sh.days:
                        no_year += sum(len(MD_NUM.findall(s)) for s in leaves(r))
                    fo.write(json.dumps(walk(r, sh), ensure_ascii=False) + "\n")
        rows.append((pid, sh.years, sh.days))
        total.update(sh.n)
        print(pid, dict(sh.n), f"numeric dates without a year: {no_year}" if sh.days else "")
    a.offsets.parent.mkdir(parents=True, exist_ok=True)
    a.offsets.write_text("patient\tyears\tdays\n" + "".join(f"{p}\t{y}\t{d}\n" for p, y, d in rows))
    print("total", dict(total))


def leaves(o):
    if isinstance(o, dict):
        for v in o.values():
            yield from leaves(v)
    elif isinstance(o, list):
        for v in o:
            yield from leaves(v)
    elif isinstance(o, str):
        yield o


if __name__ == "__main__":
    main()