solution.py

"""
solution.py
-----------
Runs the weighbridge analysis pipeline end-to-end and generates
all deliverables (output CSV, flagged data, plots, docs).

Prints a chapter-by-chapter breakdown to the terminal showing
how each insight was discovered from the raw sensor data.

Usage:
    python solution.py "path/to/assignment_wagon_data.csv"
"""

import csv
import sys
import os
from collections import Counter
from datetime import datetime

# The core algorithm lives in logic.py — import it if available,
# otherwise define everything inline so this script is self-contained.
try:
    from logic import (
        SignalDebouncer, WagonBoundaryDetector, WeighCycleTracker,
        resolve_wagon_number, clean_ocr_readings, is_valid_wagon_number,
        WeighbridgeProcessor, DEBOUNCE_THRESHOLD, AXLES_PER_WAGON,
        STABILITY_HOLD, RAKE_DEFAULT_PREFIX
    )
    IMPORTED_LOGIC = True
except ImportError:
    IMPORTED_LOGIC = False

# matplotlib is optional — plots are skipped if it's missing
try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    import numpy as np
    CAN_PLOT = True
except ImportError:
    CAN_PLOT = False
    print("  [!] matplotlib not found — skipping plot generation")
    print("      Install with: pip install matplotlib numpy")


# ──────────────────────────────────────────────────────────────
# Inline fallback: if logic.py isn't importable, define core
# components here so solution.py works standalone too.
# ──────────────────────────────────────────────────────────────
if not IMPORTED_LOGIC:
    DEBOUNCE_THRESHOLD = 6
    AXLES_PER_WAGON = 4
    STABILITY_HOLD = 15
    PREFIX_CONFIRM_RATIO = 0.70
    TIE_READ_MULTIPLIER = 2.0
    RAKE_DEFAULT_PREFIX = '224422'

    def compute_check_digit(first_ten):
        d = [int(ch) for ch in first_ten]
        odd = d[1] + d[3] + d[5] + d[7] + d[9]
        even = d[0] + d[2] + d[4] + d[6] + d[8]
        return (10 - (3 * odd + even) % 10) % 10

    def is_valid_wagon_number(number):
        if len(number) != 11 or not number.isdigit():
            return False
        return int(number[10]) == compute_check_digit(number[:10])

    class SignalDebouncer:
        def __init__(self, threshold, allowed_states=None):
            self.threshold = threshold
            self.allowed_states = allowed_states
            self.confirmed = None
            self.pending = None
            self.count = 0

        def update(self, value):
            if self.allowed_states and value not in self.allowed_states:
                return self.confirmed, False
            changed = False
            if value == self.confirmed:
                self.pending = None
                self.count = 0
            elif value == self.pending:
                self.count += 1
                if self.count >= self.threshold:
                    self.confirmed = self.pending
                    self.pending = None
                    self.count = 0
                    changed = True
            else:
                self.pending = value
                self.count = 1
            return self.confirmed, changed

    class WagonBoundaryDetector:
        def __init__(self):
            states = {'single_wagon', 'wagon_switch', 'no_wagon'}
            self.cam1 = SignalDebouncer(DEBOUNCE_THRESHOLD, states)
            self.cam2 = SignalDebouncer(DEBOUNCE_THRESHOLD, states)
            self.prev = None

        def update(self, ev1, ev2):
            state, changed = self.cam1.update(ev1)
            self.cam2.update(ev2)
            if not changed:
                return None
            old = self.prev
            self.prev = state
            if old == 'single_wagon' and state in ('wagon_switch', 'no_wagon'):
                return 'EXITED'
            if old in ('wagon_switch', 'no_wagon', None) and state == 'single_wagon':
                return 'ENTERED'
            return None

    class WeighCycleTracker:
        def __init__(self):
            self.last_status = None
            self.cycle_count = 0
            self.display = None
            self.hold = 0
            self.last_reg = None
            self.last_stable = None

        def update(self, raw_status, raw_weight):
            registered = None
            status = str(raw_status).strip()
            try:
                weight = float(raw_weight)
                if weight != weight:
                    weight = None
            except (ValueError, TypeError):
                weight = None

            if status in ('weighing---', 'weighing_over'):
                if (self.last_status and status != self.last_status
                        and self.last_status == 'weighing---'):
                    self.cycle_count += 1
                    if self.cycle_count % AXLES_PER_WAGON == 0:
                        if weight and weight > 0:
                            registered = weight
                            self.last_reg = weight
                self.last_status = status

            if weight and weight > 0:
                if weight == self.display:
                    self.hold += 1
                else:
                    self.display = weight
                    self.hold = 1
                if self.hold == STABILITY_HOLD:
                    self.last_stable = weight

            return registered

        def get_fallback(self):
            if self.last_stable and self.last_stable != self.last_reg:
                return self.last_stable
            return None

    def clean_ocr_readings(raw_list):
        cleaned = []
        for r in raw_list:
            text = str(r).strip()
            if text in ('', 'nan', 'None'):
                continue
            if '.' in text:
                text = text.split('.')[0]
            if text.isdigit() and len(text) >= 3:
                cleaned.append(text)
        return cleaned

    def resolve_wagon_number(cam1_raw, cam2_raw):
        cam1 = clean_ocr_readings(cam1_raw)
        cam2 = clean_ocr_readings(cam2_raw)
        if not cam1 and not cam2:
            return None, {}

        full1 = [r for r in cam1 if len(r) == 11]
        full2 = [r for r in cam2 if len(r) == 11]

        # Suffix
        suf1 = Counter([r[6:] for r in full1])
        suf2 = Counter([r[6:] for r in full2])
        top1 = suf1.most_common(1)[0] if suf1 else ('?', 0)
        top2 = suf2.most_common(1)[0] if suf2 else ('?', 0)
        con1 = top1[1] / len(full1) if full1 else 0
        con2 = top2[1] / len(full2) if full2 else 0

        if con1 > con2:
            suffix, reliable, side = top1[0], full1, 'CAM1'
        elif con2 > con1:
            suffix, reliable, side = top2[0], full2, 'CAM2'
        elif len(full2) > TIE_READ_MULTIPLIER * len(full1) and full2:
            suffix, reliable, side = top2[0], full2, 'CAM2'
        elif len(full1) > TIE_READ_MULTIPLIER * len(full2) and full1:
            suffix, reliable, side = top1[0], full1, 'CAM1'
        else:
            pc1 = Counter([r[:6] for r in full1]).most_common(1)[0][1] / len(full1) if full1 else 0
            pc2 = Counter([r[:6] for r in full2]).most_common(1)[0][1] / len(full2) if full2 else 0
            if pc1 >= pc2 and full1:
                suffix, reliable, side = top1[0], full1, 'CAM1'
            elif full2:
                suffix, reliable, side = top2[0], full2, 'CAM2'
            else:
                partials = [r for r in cam1 + cam2 if len(r) >= 5]
                if partials:
                    suffix = Counter([r[-5:] for r in partials]).most_common(1)[0][0]
                    reliable, side = [], 'PARTIAL'
                else:
                    return None, {}

        # Prefix
        f1 = sum(1 for r in full1 if r[:6] == RAKE_DEFAULT_PREFIX) / len(full1) if full1 else 0
        f2 = sum(1 for r in full2 if r[:6] == RAKE_DEFAULT_PREFIX) / len(full2) if full2 else 0

        if f1 >= PREFIX_CONFIRM_RATIO or f2 >= PREFIX_CONFIRM_RATIO:
            prefix, method = RAKE_DEFAULT_PREFIX, 'DEFAULT'
        elif full1 + full2 and sum(1 for r in full1 + full2 if r[:6] == '860411') / len(full1 + full2) > 0.5:
            prefix, method = '860411', 'GUARD_VAN'
        elif reliable:
            cam_pfx = Counter([r[:6] for r in reliable]).most_common(1)[0][0]
            if len(cam_pfx) == 6 and cam_pfx[:2] == '22' and cam_pfx[4:6] == '22' and cam_pfx != RAKE_DEFAULT_PREFIX:
                prefix, method = cam_pfx, 'DIFF_ALLOTMENT'
            elif cam_pfx == RAKE_DEFAULT_PREFIX:
                prefix, method = RAKE_DEFAULT_PREFIX, 'DEFAULT'
            else:
                prefix, method = RAKE_DEFAULT_PREFIX, 'DEFAULT'
        else:
            prefix, method = RAKE_DEFAULT_PREFIX, 'DEFAULT'

        info = {'prefix': prefix, 'suffix': suffix, 'method': method, 'side': side,
                'cam1_reads': len(full1), 'cam2_reads': len(full2)}
        return prefix + suffix, info

    class WeighbridgeProcessor:
        def __init__(self):
            self.boundary = WagonBoundaryDetector()
            self.weigher = WeighCycleTracker()
            self.buf1 = []
            self.buf2 = []
            self.wagon_active = False
            self.first_ocr = False
            self.eot_done = False
            self.num_queue = []
            self.wt_queue = []
            self.results = []
            self.wagon_info = []
            self.row_count = 0

        def process_row(self, row):
            self.row_count += 1
            ev1 = row.get('wagon_event', '').strip()
            ev2 = row.get('wagon_event_secondary', '').strip()
            ws = row.get('weight_status', '').strip()
            wv = row.get('weight', '')
            o1 = row.get('ocr_detection_side1', '').strip()
            o2 = row.get('ocr_detection_side2', '').strip()

            boundary = self.boundary.update(ev1, ev2)
            if boundary == 'EXITED':
                self._close_wagon()
            elif boundary == 'ENTERED':
                self.wagon_active = True
                self.buf1 = []
                self.buf2 = []

            if self.wagon_active:
                if o1 and o1 not in ('', 'nan', 'None'):
                    self.buf1.append(o1)
                if o2 and o2 not in ('', 'nan', 'None'):
                    self.buf2.append(o2)

            reg = self.weigher.update(ws, wv)
            if reg is not None:
                self.wt_queue.append(reg)
                self._pair()
            if ws == 'end_of_train':
                self._end_of_train()

        def _close_wagon(self):
            self.wagon_active = False
            number, info = resolve_wagon_number(self.buf1, self.buf2)
            self.buf1 = []
            self.buf2 = []
            if number:
                self.first_ocr = True
                self.num_queue.append(number)
                self.wagon_info.append(info)
                self._pair()
            elif not self.first_ocr:
                self.wagon_info.append({'method': 'ENGINE'})

        def _pair(self):
            while self.num_queue and self.wt_queue:
                self.results.append((self.num_queue.pop(0), self.wt_queue.pop(0)))

        def _end_of_train(self):
            if self.eot_done:
                return
            self.eot_done = True
            if self.buf1 or self.buf2:
                self._close_wagon()
            fb = self.weigher.get_fallback()
            if fb and self.num_queue:
                self.wt_queue.append(fb)
                self._pair()

        def finalize(self):
            if self.buf1 or self.buf2:
                self._close_wagon()
            fb = self.weigher.get_fallback()
            if fb and self.num_queue:
                self.wt_queue.append(fb)
                self._pair()
            return self.results


# ──────────────────────────────────────────────────────────────
# Helper: save a matplotlib figure safely
# ──────────────────────────────────────────────────────────────
def save_plot(filename, fig_dir):
    if not CAN_PLOT:
        return False
    plt.tight_layout()
    path = os.path.join(fig_dir, filename)
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    return True


# ──────────────────────────────────────────────────────────────
# Helper: detect segments from raw rows (for analysis only)
# ──────────────────────────────────────────────────────────────
def find_segments(rows):
    """Run boundary detection on raw rows to get (start, end) tuples."""
    detector = WagonBoundaryDetector()
    segments = []
    entered_at = None

    for i, row in enumerate(rows):
        ev1 = row.get('wagon_event', '').strip()
        ev2 = row.get('wagon_event_secondary', '').strip()
        event = detector.update(ev1, ev2)

        if event == 'ENTERED':
            entered_at = i
        elif event == 'EXITED' and entered_at is not None:
            segments.append((entered_at, i))
            entered_at = None

    return segments


def segment_has_ocr(rows, start, end):
    """Check if any row in the segment has OCR data."""
    for j in range(start, min(end + 1, len(rows))):
        for col in ('ocr_detection_side1', 'ocr_detection_side2'):
            val = str(rows[j].get(col, '')).strip()
            if val and val not in ('', 'nan', 'None'):
                return True
    return False


def collect_all_11digit(rows):
    """Grab every 11-digit OCR reading from the entire dataset."""
    readings = []
    for row in rows:
        for col in ('ocr_detection_side1', 'ocr_detection_side2'):
            val = str(row.get(col, '')).strip()
            if val in ('', 'nan', 'None'):
                continue
            if '.' in val:
                val = val.split('.')[0]
            if val.isdigit() and len(val) == 11:
                readings.append(val)
    return readings


def per_wagon_camera_stats(rows, segments, segment_types):
    """For each OCR segment, compute suffix consistency per camera."""
    stats = []
    for i, (start, end) in enumerate(segments):
        if segment_types[i] == 'ENGINE':
            continue

        cam1_full = []
        cam2_full = []
        for j in range(start, min(end + 1, len(rows))):
            for col, bucket in [('ocr_detection_side1', cam1_full), ('ocr_detection_side2', cam2_full)]:
                val = str(rows[j].get(col, '')).strip()
                if val in ('', 'nan', 'None'):
                    continue
                if '.' in val:
                    val = val.split('.')[0]
                if val.isdigit() and len(val) == 11:
                    bucket.append(val)

        suf1 = Counter([r[6:] for r in cam1_full])
        suf2 = Counter([r[6:] for r in cam2_full])
        con1 = suf1.most_common(1)[0][1] / len(cam1_full) if cam1_full else 0
        con2 = suf2.most_common(1)[0][1] / len(cam2_full) if cam2_full else 0

        stats.append({
            'cam1_consistency': con1,
            'cam2_consistency': con2,
            'cam1_reads': len(cam1_full),
            'cam2_reads': len(cam2_full),
        })
    return stats


# ──────────────────────────────────────────────────────────────
# MAIN: the full analysis with terminal output + plots
# ──────────────────────────────────────────────────────────────
def main(input_csv):
    output_dir = 'output'
    fig_dir = os.path.join(output_dir, 'figures')
    os.makedirs(fig_dir, exist_ok=True)

    # Load data
    with open(input_csv, 'r') as f:
        rows = list(csv.DictReader(f))
    total_rows = len(rows)

    print(f"""
{'='*72}
  REAL-TIME TRAIN WEIGHBRIDGE ANALYSIS
  A Data Engineering Story
{'='*72}
  Input : {input_csv}
  Rows  : {total_rows:,}
  Time  : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
{'='*72}""")

    # ── Run the processor ──
    processor = WeighbridgeProcessor()
    for row in rows:
        processor.process_row(row)
    final_output = processor.finalize()

    # ── Gather analysis data ──
    segments = find_segments(rows)
    segment_types = []
    for start, end in segments:
        has_ocr = segment_has_ocr(rows, start, end)
        segment_types.append('OCR' if has_ocr else 'ENGINE')

    engine_count = segment_types.count('ENGINE')
    ocr_count = segment_types.count('OCR')
    cargo_count = ocr_count - 1  # last OCR segment is guard van
    guard_count = 1

    # Weight registrations (separate tracker for analysis)
    analysis_weigher = WeighCycleTracker()
    weight_registrations = []
    for i, row in enumerate(rows):
        reg = analysis_weigher.update(row.get('weight_status', ''), row.get('weight', ''))
        if reg is not None:
            weight_registrations.append((i, reg))

    zero_weights = [(i, w) for i, w in weight_registrations if w == 0]
    cargo_weights = [(i, w) for i, w in weight_registrations if w > 0]
    guard_van_weight = final_output[-1][1] if final_output else 0

    # All 11-digit OCR reads
    all_11digit = collect_all_11digit(rows)

    # Per-wagon camera stats
    camera_stats = per_wagon_camera_stats(rows, segments, segment_types)

    # Enriched results
    enriched = []
    for i, (wagon_num, weight) in enumerate(final_output):
        info_idx = i + engine_count
        info = processor.wagon_info[info_idx] if info_idx < len(processor.wagon_info) else {}
        prefix = info.get('prefix', wagon_num[:6] if len(wagon_num) >= 6 else '?')
        suffix = info.get('suffix', wagon_num[6:] if len(wagon_num) >= 11 else '?')
        method = info.get('method', 'DEFAULT')
        check_ok = is_valid_wagon_number(wagon_num)
        enriched.append({
            'number': wagon_num, 'weight': weight, 'prefix': prefix,
            'suffix': suffix, 'method': method, 'check_ok': check_ok,
        })

    # ─────────────────────────────────────────────────
    # CHAPTER 1: Train Structure
    # ─────────────────────────────────────────────────
    print(f"""
{'─'*72}
  CHAPTER 1: WHAT'S IN THE DATA?
{'─'*72}

  Debounced wagon_event ({DEBOUNCE_THRESHOLD}-frame threshold) reveals:

  Total segments:  {len(segments)}
  Engine (no OCR): {engine_count}   — locomotive, different axle config
  Cargo wagons:    {cargo_count}  — loaded, OCR readable, 73-95 tonnes
  Guard van:       {guard_count}   — brake van at rear, 22.6t
""")

    if CAN_PLOT:
        fig, ax = plt.subplots(figsize=(16, 2.5))
        for i, (seg_type, (start, end)) in enumerate(zip(segment_types, segments)):
            color = '#ef4444' if seg_type == 'ENGINE' else '#3b82f6'
            if seg_type == 'OCR' and i == len(segments) - 1:
                color = '#10b981'  # guard van
            ax.barh(0, end - start, left=start, height=0.6, color=color, edgecolor='white', linewidth=0.5)
        legend_items = [
            mpatches.Patch(color='#ef4444', label=f'Engine ({engine_count})'),
            mpatches.Patch(color='#3b82f6', label=f'Cargo ({cargo_count})'),
            mpatches.Patch(color='#10b981', label='Guard Van'),
        ]
        ax.legend(handles=legend_items, loc='upper right', fontsize=9)
        ax.set_xlabel('Row index')
        ax.set_yticks([])
        ax.set_title('Train Structure')
        if save_plot('01_train_structure.png', fig_dir):
            print(f"  >> figures/01_train_structure.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 2: Weight Detection
    # ─────────────────────────────────────────────────
    min_cargo = min(w for _, w in cargo_weights)
    max_cargo = max(w for _, w in cargo_weights)

    print(f"""
{'─'*72}
  CHAPTER 2: WEIGHT DETECTION (4-cycle pattern)
{'─'*72}

  Registrations:  {len(weight_registrations)}
  Zero (engine):  {len(zero_weights)}
  Cargo weights:  {len(cargo_weights)}  ({min_cargo:.1f}t to {max_cargo:.1f}t)
  Guard van:      {guard_van_weight}t  (stability fallback, not 4-cycle)
""")

    if CAN_PLOT:
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

        # Histogram
        wt_values = [w for _, w in cargo_weights]
        bins = [70, 75, 80, 85, 90, 95, 100]
        bin_counts = [len([w for w in wt_values if bins[i] <= w < bins[i+1]]) for i in range(len(bins)-1)]
        ax1.bar([f'{bins[i]}-{bins[i+1]}t' for i in range(len(bins)-1)], bin_counts,
                color='#3b82f6', alpha=0.8, edgecolor='white', linewidth=0.5)
        ax1.set_ylabel('Count')
        ax1.set_title('Cargo Weight Distribution')

        # Pie
        ax2.pie([len(cargo_weights), 1, engine_count],
                labels=[f'4-Cycle Cargo ({len(cargo_weights)})', f'Guard Van (1)', f'Engine ({engine_count})'],
                colors=['#3b82f6', '#10b981', '#ef4444'],
                autopct='%1.0f%%', startangle=90)
        ax2.set_title('Detection Method')

        if save_plot('02_weight_detection.png', fig_dir):
            print(f"  >> figures/02_weight_detection.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 3: OCR Number Structure
    # ─────────────────────────────────────────────────
    print(f"""
{'─'*72}
  CHAPTER 3: CRACKING THE WAGON NUMBER FORMAT
{'─'*72}

  {len(all_11digit):,} eleven-digit OCR readings analyzed digit-by-digit:
""")
    print(f"  {'Pos':>5} {'Top':>6} {'Count':>7} {'Pct':>7} {'Type':>12}")
    print(f"  {'─'*42}")

    digit_uniqueness = []
    for pos in range(11):
        digits = [r[pos] for r in all_11digit]
        top_digit, top_count = Counter(digits).most_common(1)[0]
        pct = 100 * top_count / len(digits)
        kind = "FIXED" if pct > 90 else "SEMI" if pct > 60 else "VARIABLE"
        unique_count = len(set(digits))
        digit_uniqueness.append(unique_count)
        print(f"  C{pos+1:>3}    '{top_digit}'  {top_count:>6}  {pct:>5.1f}%  {kind:>12}")

    print(f"""
  Positions 1-6 = PREFIX (constrained) — shared across the rake
  Positions 7-11 = SUFFIX (variable) — unique per wagon
""")

    if CAN_PLOT:
        fig, ax = plt.subplots(figsize=(12, 5))
        bar_colors = ['#10b981' if u <= 6 else '#ef4444' for u in digit_uniqueness]
        ax.bar(range(11), digit_uniqueness, color=bar_colors, alpha=0.8, edgecolor='white')
        ax.axvline(5.5, color='#6366f1', linewidth=2, linestyle='--', alpha=0.7)
        ax.text(2.5, max(digit_uniqueness) * 0.85, 'PREFIX', ha='center',
                fontsize=12, color='#10b981', fontweight='bold')
        ax.text(8, max(digit_uniqueness) * 0.85, 'SUFFIX', ha='center',
                fontsize=12, color='#ef4444', fontweight='bold')
        ax.set_xticks(range(11))
        ax.set_xticklabels([f'C{i+1}' for i in range(11)])
        ax.set_ylabel('Unique digits observed')
        ax.set_title('Digit Variability per Position')
        for i, u in enumerate(digit_uniqueness):
            ax.text(i, u + 0.15, str(u), ha='center', fontsize=9, fontweight='bold')
        if save_plot('03_digit_variability.png', fig_dir):
            print(f"  >> figures/03_digit_variability.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 4: Dominant Prefix
    # ─────────────────────────────────────────────────
    prefix_counts = Counter([r[:6] for r in all_11digit])
    dominant_pct = 100 * prefix_counts[RAKE_DEFAULT_PREFIX] / len(all_11digit)

    print(f"""
{'─'*72}
  CHAPTER 4: THE DOMINANT PREFIX — {RAKE_DEFAULT_PREFIX}
{'─'*72}
""")
    for pfx, cnt in prefix_counts.most_common(8):
        pct = 100 * cnt / len(all_11digit)
        bar = '#' * int(pct / 2)
        print(f"    {pfx}: {cnt:>4} ({pct:>5.1f}%) {bar}")

    print(f"""
  '{RAKE_DEFAULT_PREFIX}' = BOXN-HL type, owner zone 44, year 2022
  Dominates at {dominant_pct:.0f}% of all readings
""")

    if CAN_PLOT:
        top_pfx = prefix_counts.most_common(10)
        fig, ax = plt.subplots(figsize=(12, 5))
        pfx_labels = [p for p, _ in top_pfx]
        pfx_values = [c for _, c in top_pfx]
        pfx_colors = ['#3b82f6' if p == RAKE_DEFAULT_PREFIX else '#f59e0b' if p == '860411' else '#d1d5db'
                      for p in pfx_labels]
        ax.barh(range(len(pfx_labels)), pfx_values, color=pfx_colors, edgecolor='white')
        ax.set_yticks(range(len(pfx_labels)))
        ax.set_yticklabels(pfx_labels, fontfamily='monospace')
        ax.invert_yaxis()
        ax.set_xlabel('Count')
        ax.set_title('Prefix Frequency')
        for i, v in enumerate(pfx_values):
            ax.text(v + 2, i, str(v), va='center', fontsize=9)
        if save_plot('04_prefix_distribution.png', fig_dir):
            print(f"  >> figures/04_prefix_distribution.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 5: Camera Consistency
    # ─────────────────────────────────────────────────
    cam1_wins = sum(1 for s in camera_stats if s['cam1_consistency'] > s['cam2_consistency'])
    cam2_wins = sum(1 for s in camera_stats if s['cam2_consistency'] > s['cam1_consistency'])
    ties = sum(1 for s in camera_stats if s['cam1_consistency'] == s['cam2_consistency'])

    print(f"""
{'─'*72}
  CHAPTER 5: WHICH CAMERA READS BETTER?
{'─'*72}

  Camera 1 wins: {cam1_wins}
  Camera 2 wins: {cam2_wins}
  Tie:           {ties}

  Suffix comes from the more consistent camera.
  Tie-breaker: >2x reads, then prefix consistency.
""")

    if CAN_PLOT:
        fig, ax = plt.subplots(figsize=(16, 5))
        x_axis = np.arange(len(camera_stats))
        c1_vals = [s['cam1_consistency'] for s in camera_stats]
        c2_vals = [s['cam2_consistency'] for s in camera_stats]
        ax.bar(x_axis - 0.15, c1_vals, 0.3, color='#8b5cf6', alpha=0.7, label='Camera 1')
        ax.bar(x_axis + 0.15, c2_vals, 0.3, color='#ec4899', alpha=0.7, label='Camera 2')
        ax.set_xlabel('Wagon #')
        ax.set_ylabel('Suffix Consistency')
        ax.set_title('Camera Consistency per Wagon')
        ax.legend()
        ax.set_ylim(0, 1.1)
        if save_plot('05_camera_consistency.png', fig_dir):
            print(f"  >> figures/05_camera_consistency.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 6: Different Allotment
    # ─────────────────────────────────────────────────
    diff_wagons = [(i + 1, r) for i, r in enumerate(enriched) if 'DIFF' in r['method']]
    allotment_zones = Counter(r['prefix'][2:4] for _, r in diff_wagons)

    print(f"""
{'─'*72}
  CHAPTER 6: DIFFERENT ALLOTMENT WAGONS
{'─'*72}

  {len(diff_wagons)} wagons from other railway zones:
""")
    for wagon_num, record in diff_wagons:
        print(f"    W{wagon_num:>2}: {record['number']}  zone {record['prefix'][2:4]}  ({record['prefix']})")

    print()
    for zone, count in allotment_zones.most_common():
        print(f"    22-{zone}-22: {count} wagon(s)")

    if CAN_PLOT:
        fig, ax = plt.subplots(figsize=(18, 2.5))
        for i, record in enumerate(enriched):
            method = record['method']
            color = '#3b82f6' if 'DEFAULT' in method else '#ef4444' if method == 'GUARD_VAN' else '#f59e0b'
            ax.bar(i, 1, color=color, edgecolor='white', linewidth=0.5)
            if 'DIFF' in method or method == 'GUARD_VAN':
                ax.text(i, 0.5, record['prefix'], ha='center', va='center',
                        fontsize=5, rotation=90, color='white', fontweight='bold')
        legend_items = [
            mpatches.Patch(color='#3b82f6', label=f'224422 ({sum(1 for r in enriched if "DEFAULT" in r["method"])})'),
            mpatches.Patch(color='#f59e0b', label=f'Different allotment ({len(diff_wagons)})'),
            mpatches.Patch(color='#ef4444', label='Guard van'),
        ]
        ax.legend(handles=legend_items, loc='upper right', fontsize=9)
        ax.set_xlabel('Wagon #')
        ax.set_yticks([])
        ax.set_title('Prefix Map')
        if save_plot('06_prefix_map.png', fig_dir):
            print(f"\n  >> figures/06_prefix_map.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 7: Serial Pattern
    # ─────────────────────────────────────────────────
    serials = []
    for record in enriched:
        if record['prefix'] == '860411':
            continue
        try:
            serials.append(int(record['number'][6:10]))
        except (ValueError, IndexError):
            pass

    in_range = [s for s in serials if 1000 <= s <= 1060]
    out_of_range = [s for s in serials if s < 1000 or s > 1060]
    missing_serials = sorted(set(range(min(in_range), max(in_range) + 1)) - set(in_range)) if in_range else []

    print(f"""
{'─'*72}
  CHAPTER 7: SERIAL NUMBER PATTERN
{'─'*72}

  In range 1000-1060: {len(in_range)} wagons
  Out of range:       {len(out_of_range)}
  Missing:            {missing_serials}

  Near-consecutive batch validates the OCR resolution.
""")

    if CAN_PLOT:
        fig, ax = plt.subplots(figsize=(16, 5))
        for i, serial in enumerate(serials):
            color = '#3b82f6' if 1000 <= serial <= 1060 else '#ef4444'
            ax.bar(i + 1, serial, color=color, edgecolor='white', linewidth=0.5)
        ax.axhline(1000, color='#10b981', linestyle='--', alpha=0.5)
        ax.axhline(1060, color='#10b981', linestyle='--', alpha=0.5)
        ax.set_xlabel('Wagon #')
        ax.set_ylabel('Serial Number')
        ax.set_title('Serial Numbers — Consecutive Batch 1000-1060')
        if save_plot('07_serial_pattern.png', fig_dir):
            print(f"  >> figures/07_serial_pattern.png")

    # ─────────────────────────────────────────────────
    # CHAPTER 8: Final Output Table
    # ─────────────────────────────────────────────────
    check_pass_count = sum(1 for r in enriched if r['check_ok'])
    default_count = sum(1 for r in enriched if 'DEFAULT' in r['method'])
    diff_count = len(diff_wagons)

    print(f"""
{'='*72}
  CHAPTER 8: THE FINAL ANSWER
{'='*72}

  {'W#':>3} {'Wagon Number':>13} {'Weight':>7}  {'Prefix':>8} {'Method':>18} {'Check':>6}
  {'─'*62}""")

    for i, record in enumerate(enriched):
        check_mark = 'ok' if record['check_ok'] else '--'
        flag = ' *' if 'DIFF' in record['method'] else (' #' if record['method'] == 'GUARD_VAN' else '')
        print(f"  {i+1:>3} {record['number']:>13} {record['weight']:>7}  "
              f"{record['prefix']:>8} {record['method']:>18} {check_mark:>6}{flag}")

    if CAN_PLOT:
        fig, axes = plt.subplots(3, 1, figsize=(18, 10))

        # Prefix type
        for i, record in enumerate(enriched):
            method = record['method']
            color = '#3b82f6' if 'DEFAULT' in method else '#ef4444' if method == 'GUARD_VAN' else '#f59e0b'
            axes[0].bar(i, 1, color=color, edgecolor='white')
        axes[0].set_title('Prefix type (blue=default, amber=diff allotment, red=guard van)')
        axes[0].set_yticks([])

        # Check digit
        for i, record in enumerate(enriched):
            color = '#10b981' if record['check_ok'] else '#ef4444'
            axes[1].bar(i, 1, color=color, edgecolor='white')
        axes[1].set_title('Check digit validation (green=pass, red=fail)')
        axes[1].set_yticks([])

        # Weights
        for i, record in enumerate(enriched):
            color = '#3b82f6' if record['weight'] > 50 else '#f59e0b'
            axes[2].bar(i, record['weight'], color=color, edgecolor='white')
        axes[2].set_ylabel('Weight (t)')
        axes[2].set_xlabel('Wagon #')
        axes[2].set_title('Registered weights')

        if save_plot('08_final_pipeline.png', fig_dir):
            print(f"\n  >> figures/08_final_pipeline.png")

    # ─────────────────────────────────────────────────
    # WRITE DELIVERABLES
    # ─────────────────────────────────────────────────
    print(f"""
{'─'*72}
  WRITING DELIVERABLES
{'─'*72}""")

    # output.csv
    output_csv_path = os.path.join(output_dir, 'output.csv')
    with open(output_csv_path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['wagon_number', 'weight'])
        for wagon_num, weight in final_output:
            writer.writerow([wagon_num, weight])
    print(f"  >> output/output.csv          ({len(final_output)} rows)")

    # flagged_data.csv — input with wagon IDs and allotment flags
    flagged_path = os.path.join(output_dir, 'flagged_data.csv')
    original_headers = list(csv.DictReader(open(input_csv)).fieldnames)
    extra_headers = ['wagon_id', 'prefix', 'allotment_flag']

    # Build per-row flags
    row_wagon_id = [''] * total_rows
    row_prefix = [''] * total_rows
    row_flag = [''] * total_rows

    wagon_index = 0
    for seg_i, (start, end) in enumerate(segments):
        if segment_types[seg_i] == 'ENGINE':
            continue
        if wagon_index >= len(enriched):
            break

        record = enriched[wagon_index]
        for j in range(start, min(end + 1, total_rows)):
            row_wagon_id[j] = f'W{wagon_index + 1}'
            row_prefix[j] = record['prefix']
            if 'DIFF' in record['method']:
                row_flag[j] = f"DIFF_{record['prefix']}"
            elif record['method'] == 'GUARD_VAN':
                row_flag[j] = 'GUARD_VAN'
        wagon_index += 1

    with open(flagged_path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(original_headers + extra_headers)
        with open(input_csv) as inp:
            for i, row in enumerate(csv.DictReader(inp)):
                values = [row.get(h, '') for h in original_headers]
                values.extend([row_wagon_id[i], row_prefix[i], row_flag[i]])
                writer.writerow(values)
    print(f"  >> output/flagged_data.csv     (input + allotment flags)")



    # ─────────────────────────────────────────────────
    # Summary
    # ─────────────────────────────────────────────────
    print(f"""
{'='*72}
  DONE
{'='*72}

  Train:   [{engine_count} Engine] + [{cargo_count} Cargo] + [Guard Van]
  Weights: {len(cargo_weights)} cargo ({min_cargo:.1f}t-{max_cargo:.1f}t) + guard van ({guard_van_weight}t)
  OCR:     {default_count} default + {diff_count} diff allotment + {guard_count} guard van
  Check:   {check_pass_count}/{len(final_output)} pass
  Output:  {len(final_output)} rows in output/output.csv

  output/
    output.csv, flagged_data.csv, figures/ (8 plots)
{'='*72}
""")


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python solution.py <path/to/assignment_wagon_data.csv>")
        sys.exit(1)
    main(sys.argv[1])