#!/usr/bin/env python3
"""
csu_data_check.py — data quality check + EDA for CSU interval-usage exports.

Run: python3 csu_data_check.py ["path/to/folder"]

Default folder: data/CSU utility bills/CSU 15-min data 06-08_to_07-02

What it does:
  1. Reads every CSV in the folder that matches the CSU interval-export schema
     (Account Number, Meter Number, Day, Time, Consumption (kWh), kVARh).
  2. Forward-fills Account/Meter Number (CSU only populates it on the first row).
  3. Merges all files into one timeline, de-duplicating overlapping timestamps.
     When two files disagree on the same timestamp, the non-zero reading wins
     (CSU's weekly exports sometimes end in a run of erroneous 0.0000 rows that
     a follow-up "patch" export corrects) — every such resolution is logged.
  4. Flags data-quality issues: gaps in the 15-min cadence, non-15-min steps,
     duplicate timestamps, zero-reading runs, negative values.
  5. Runs a quick EDA: date range, total kWh, daily totals, peak 15-min demand,
     load factor, kVARh / estimated power factor stats.
  6. Writes a combined, deduped, time-sorted CSV + a text report into the
     folder, and prints a summary to stdout.

Re-run this any time new CSU export files are dropped into the folder.
"""

import csv
import glob
import os
import sys
from datetime import datetime, timedelta
from collections import Counter

EXPECTED_HEADER = ["Account Number", "Meter Number", "Day", "Time",
                    "Consumption (kWh)", "kVARh"]
INTERVAL_MIN = 15


def find_schema_files(folder):
    """Return CSVs in folder matching the CSU interval-export schema."""
    matches = []
    for path in sorted(glob.glob(os.path.join(folder, "*.csv"))):
        with open(path, newline="") as fh:
            rows = list(csv.reader(fh))
        for row in rows[:6]:
            if [c.strip() for c in row] == EXPECTED_HEADER:
                matches.append(path)
                break
    return matches


def parse_file(path):
    """Parse one CSU interval CSV into a list of (datetime, kwh, kvarh, account, meter)."""
    with open(path, newline="") as fh:
        rows = list(csv.reader(fh))

    header_idx = next(i for i, row in enumerate(rows)
                       if [c.strip() for c in row] == EXPECTED_HEADER)
    data_rows = rows[header_idx + 1:]

    records = []
    last_account, last_meter = None, None
    bad_rows = 0
    for row in data_rows:
        if not row or not row[2]:
            continue
        account, meter, day, time_, kwh, kvarh = (row + [""] * 6)[:6]
        last_account = account.strip() or last_account
        last_meter = meter.strip() or last_meter
        try:
            dt = datetime.strptime(f"{day.strip()} {time_.strip()}", "%m/%d/%Y %I:%M:%S %p")
            kwh_f = float(kwh)
            kvarh_f = float(kvarh) if kvarh.strip() else 0.0
        except ValueError:
            bad_rows += 1
            continue
        records.append((dt, kwh_f, kvarh_f, last_account, last_meter))
    return records, bad_rows


PATCH_FILE_MAX_ROWS = 50  # small supplemental exports are treated as authoritative corrections


def merge_records(file_records):
    """Merge records from multiple files, resolving duplicate timestamps.

    CSU occasionally issues small "patch" exports (a handful of rows) to
    correct a corrupted tail window in a full weekly export (the weekly file
    can show a decaying-to-zero run of bad readings right before/at a boundary).
    Patch files (few rows) always win over larger regular exports for any
    timestamp they both cover; conflicts between two same-size-tier files are
    left as flagged, unresolved disagreements for manual review.
    """
    by_ts = {}
    conflicts = []
    # Load largest files first so smaller "patch" files are applied last and
    # can unconditionally override.
    ordered = sorted(file_records, key=lambda fr: len(fr[1]), reverse=True)
    file_sizes = {path: len(records) for path, records in ordered}

    for path, records in ordered:
        is_patch = file_sizes[path] <= PATCH_FILE_MAX_ROWS
        for dt, kwh, kvarh, account, meter in records:
            if dt not in by_ts:
                by_ts[dt] = (kwh, kvarh, account, meter, path)
                continue
            existing = by_ts[dt]
            if existing[0] == kwh:
                continue  # identical duplicate, nothing to resolve
            existing_is_patch = file_sizes[existing[4]] <= PATCH_FILE_MAX_ROWS
            new_entry = (kwh, kvarh, account, meter, path)
            if is_patch and not existing_is_patch:
                conflicts.append((dt, existing, new_entry, "patch_override"))
                by_ts[dt] = new_entry
            elif existing_is_patch and not is_patch:
                conflicts.append((dt, existing, new_entry, "kept_existing_patch"))
            elif existing[0] == 0.0 and kwh != 0.0:
                conflicts.append((dt, existing, new_entry, "replaced_zero"))
                by_ts[dt] = new_entry
            elif kwh == 0.0:
                conflicts.append((dt, existing, new_entry, "kept_existing_over_zero"))
            else:
                conflicts.append((dt, existing, new_entry, "unresolved_disagreement"))
    return by_ts, conflicts


TRAILING_ZERO_TRIM_THRESHOLD = 3  # intervals (45 min) — CSU export/download-cutoff
                                    # artifacts consistently show a decaying-then-hard-zero
                                    # tail; real consumption doesn't sustain exact 0.0000
                                    # for this long, so trim any such run found at the very
                                    # end of the merged timeline.


def trim_trailing_zero_run(by_ts):
    """Drop a corrupted trailing run of zero readings at the end of the merged
    timeline (a recurring CSU export artifact — readings decay abnormally then
    hard-cut to 0.0000 through the download cutoff). Only trims if the zero
    run is the very last thing in the data and at least
    TRAILING_ZERO_TRIM_THRESHOLD intervals long; leaves everything else
    (including zero runs elsewhere, already fixed by patch files) untouched.
    Returns (trimmed_by_ts, trim_info_or_None).
    """
    timeline = sorted(by_ts.items())
    if not timeline:
        return by_ts, None

    trim_from_idx = None
    for i in range(len(timeline) - 1, -1, -1):
        if timeline[i][1][0] == 0.0:
            trim_from_idx = i
        else:
            break

    if trim_from_idx is None:
        return by_ts, None
    run_len = len(timeline) - trim_from_idx
    if run_len < TRAILING_ZERO_TRIM_THRESHOLD:
        return by_ts, None

    trim_start_ts = timeline[trim_from_idx][0]
    trimmed = dict(timeline[:trim_from_idx])
    info = {
        "trim_start": trim_start_ts,
        "trim_end": timeline[-1][0],
        "rows_removed": run_len,
        "new_last_ts": timeline[trim_from_idx - 1][0] if trim_from_idx > 0 else None,
    }
    return trimmed, info


def analyze(by_ts):
    timeline = sorted(by_ts.items())
    dts = [t[0] for t in timeline]
    kwhs = [v[0] for _, v in timeline]

    gaps = []
    non_standard_steps = Counter()
    for (t0, _), (t1, _) in zip(timeline, timeline[1:]):
        step_min = (t1 - t0).total_seconds() / 60
        non_standard_steps[step_min] += 1
        if step_min != INTERVAL_MIN:
            gaps.append((t0, t1, step_min))

    zero_runs = []
    run_start = None
    for (t, v) in timeline:
        if v[0] == 0.0:
            if run_start is None:
                run_start = t
        else:
            if run_start is not None:
                zero_runs.append((run_start, t))
                run_start = None
    if run_start is not None:
        zero_runs.append((run_start, None))

    negatives = [(t, v[0]) for t, v in timeline if v[0] < 0]

    daily_totals = {}
    for t, v in timeline:
        day = t.date()
        daily_totals[day] = daily_totals.get(day, 0.0) + v[0]

    peak_ts, peak_kwh = max(((t, v[0]) for t, v in timeline), key=lambda x: x[1])
    total_kwh = sum(kwhs)
    avg_kw = total_kwh / (len(timeline) * INTERVAL_MIN / 60) if timeline else 0
    peak_kw = peak_kwh / (INTERVAL_MIN / 60)
    load_factor = avg_kw / peak_kw if peak_kw else 0

    kvarh_total = sum(v[1] for _, v in timeline)
    # rough est. power factor from period totals: PF = kWh / sqrt(kWh^2 + kVARh^2)
    est_pf = total_kwh / ((total_kwh ** 2 + kvarh_total ** 2) ** 0.5) if total_kwh else 0

    return {
        "timeline": timeline,
        "n_records": len(timeline),
        "date_range": (dts[0], dts[-1]) if dts else (None, None),
        "step_counts": non_standard_steps,
        "gaps": gaps,
        "zero_runs": zero_runs,
        "negatives": negatives,
        "daily_totals": daily_totals,
        "total_kwh": total_kwh,
        "peak_ts": peak_ts,
        "peak_kwh": peak_kwh,
        "peak_kw": peak_kw,
        "avg_kw": avg_kw,
        "load_factor": load_factor,
        "kvarh_total": kvarh_total,
        "est_power_factor": est_pf,
    }


def write_outputs(folder, stats, conflicts, per_file_bad_rows, trim_info=None):
    combined_path = os.path.join(folder, "combined_clean.csv")
    with open(combined_path, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["Timestamp", "Day", "Time", "Consumption (kWh)", "kVARh", "Account", "Meter", "SourceFile"])
        for t, (kwh, kvarh, account, meter, path) in stats["timeline"]:
            w.writerow([t.isoformat(), t.strftime("%m/%d/%Y"), t.strftime("%I:%M:%S %p"),
                        kwh, kvarh, account, meter, os.path.basename(path)])

    report_path = os.path.join(folder, "data_check_report.txt")
    lines = []
    lines.append(f"CSU 15-min Data Check Report — generated {datetime.now().isoformat(timespec='seconds')}")
    lines.append(f"Folder: {folder}")
    lines.append("")
    if trim_info:
        lines.append(f"TRIMMED: {trim_info['rows_removed']} rows of corrupted trailing zero-readings "
                      f"removed ({trim_info['trim_start']} to {trim_info['trim_end']}) — export/"
                      f"download-cutoff artifact, not real consumption. New last reliable reading: "
                      f"{trim_info['new_last_ts']}")
        lines.append("")
    lines.append(f"Records after merge/dedup: {stats['n_records']}")
    lines.append(f"Date range: {stats['date_range'][0]} to {stats['date_range'][1]}")
    lines.append("")
    lines.append("Interval step distribution (minutes: count):")
    for step, count in sorted(stats["step_counts"].items()):
        flag = "" if step == INTERVAL_MIN else "  <-- NON-STANDARD"
        lines.append(f"  {step:>6}: {count}{flag}")
    lines.append("")
    if stats["gaps"]:
        lines.append(f"Gaps / non-15-min steps found: {len(stats['gaps'])}")
        for t0, t1, step in stats["gaps"]:
            lines.append(f"  {t0} -> {t1}  ({step:.0f} min gap)")
    else:
        lines.append("No gaps found — full 15-min cadence maintained.")
    lines.append("")
    if stats["zero_runs"]:
        lines.append(f"Zero-reading runs remaining after merge: {len(stats['zero_runs'])}")
        for start, end in stats["zero_runs"]:
            lines.append(f"  {start} -> {end if end else '(end of data)'}")
    else:
        lines.append("No zero-reading runs remaining after merge.")
    lines.append("")
    if stats["negatives"]:
        lines.append(f"Negative consumption values found: {len(stats['negatives'])}")
        for t, v in stats["negatives"]:
            lines.append(f"  {t}: {v}")
    else:
        lines.append("No negative consumption values found.")
    lines.append("")
    lines.append(f"Duplicate-timestamp conflicts resolved during merge: {len(conflicts)}")
    for dt, existing, new, reason in conflicts:
        lines.append(f"  {dt}: existing={existing[0]} ({os.path.basename(existing[4])}) "
                      f"vs new={new[0]} ({os.path.basename(new[4])}) -> {reason}")
    lines.append("")
    if any(per_file_bad_rows.values()):
        lines.append("Unparseable rows per file:")
        for path, n in per_file_bad_rows.items():
            if n:
                lines.append(f"  {os.path.basename(path)}: {n} bad rows skipped")
        lines.append("")

    lines.append("--- EDA summary ---")
    lines.append(f"Total consumption: {stats['total_kwh']:.2f} kWh")
    lines.append(f"Average demand: {stats['avg_kw']:.2f} kW")
    lines.append(f"Peak 15-min demand: {stats['peak_kw']:.2f} kW at {stats['peak_ts']} ({stats['peak_kwh']:.2f} kWh interval)")
    lines.append(f"Load factor (avg/peak): {stats['load_factor']:.2%}")
    lines.append(f"Total kVARh: {stats['kvarh_total']:.2f}")
    lines.append(f"Estimated overall power factor: {stats['est_power_factor']:.3f}")
    lines.append("")
    lines.append("Daily totals (kWh):")
    for day, total in sorted(stats["daily_totals"].items()):
        lines.append(f"  {day}: {total:.2f}")

    report_text = "\n".join(lines)
    with open(report_path, "w") as fh:
        fh.write(report_text + "\n")

    return combined_path, report_path, report_text


def main():
    default_folder = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "data", "CSU utility bills", "CSU 15-min data 06-08_to_07-02",
    )
    folder = sys.argv[1] if len(sys.argv) > 1 else default_folder
    if not os.path.isdir(folder):
        print(f"Folder not found: {folder}")
        sys.exit(1)

    files = find_schema_files(folder)
    if not files:
        print(f"No CSU interval-export CSVs found in {folder}")
        sys.exit(1)

    print(f"Found {len(files)} matching CSV file(s) in {folder}:")
    file_records = []
    per_file_bad_rows = {}
    for path in files:
        records, bad_rows = parse_file(path)
        per_file_bad_rows[path] = bad_rows
        print(f"  {os.path.basename(path)}: {len(records)} rows"
              + (f", {bad_rows} unparseable" if bad_rows else ""))
        file_records.append((path, records))

    by_ts, conflicts = merge_records(file_records)
    by_ts, trim_info = trim_trailing_zero_run(by_ts)
    if trim_info:
        print(f"\nTrimmed corrupted trailing zero-run: {trim_info['rows_removed']} rows from "
              f"{trim_info['trim_start']} to {trim_info['trim_end']} (export/download-cutoff artifact). "
              f"New last reliable reading: {trim_info['new_last_ts']}")
    stats = analyze(by_ts)
    combined_path, report_path, report_text = write_outputs(
        folder, stats, conflicts, per_file_bad_rows, trim_info)

    print()
    print(report_text)
    print()
    print(f"Combined clean CSV written to: {combined_path}")
    print(f"Report written to: {report_path}")


if __name__ == "__main__":
    main()