"""LearnAI P02 — inspect local Iris arrays and plot development evidence.
Python 3.11+; install studio-requirements.txt in an isolated environment.
python array-inspection.py --csv iris.csv --out inspection-01
python array-inspection.py --check  # Learner TODO intentionally fails initially.
python array-inspection.py --csv iris.csv --out inspection-02 --learner

No model is trained. Training/development rows are public repeated evidence;
no final evaluation score or evaluation scatter is produced. The CSV contains
all labels and is public, so this is not secret or previously unseen data.
Dataset: Fisher, R. (1936). Iris. UCI, https://doi.org/10.24432/C56C76,
CC BY 4.0. Generated row IDs are not specimen identities. Equal measurement
vectors form groups; this is a practical duplicate rule, not proven identity.
"""
import argparse
import csv
import hashlib
import io
import json
import platform
import sys
import unittest
from pathlib import Path
import numpy as np

CSV_HASH = 'b34e5958ea7f7625f6a43c060b28d9848e6b810c5d09652b2335ef4619db9df8'
FEATURES = ['sepal_length_cm', 'sepal_width_cm', 'petal_length_cm', 'petal_width_cm']
SPECIES = ['Iris-setosa', 'Iris-versicolor', 'Iris-virginica']
PROTOCOL = 'iris-array-inspection-grouped-v1'


def learner_standardize(training, other):
    """Return (training_z, other_z, mean, scale), as NEW float64 arrays.
    Inputs: finite 2D arrays, 1..1000 rows and 1..16 matching feature columns;
    absolute input values <= 1e6. Reject violations with ValueError. Do not
    mutate inputs. Fit mean and population SD PER COLUMN on training only.
    Replace zero SD by 1. Return mean/scale with shape (features,), both z
    arrays with original input shapes. No missing-value policy in this task:
    nonfinite data must be repaired upstream, not silently accepted here.
    """
    raise NotImplementedError('Implement learner_standardize and run --check.')


def validate_arrays(training, other):
    a, b = np.array(training, dtype=np.float64, copy=True), np.array(other, dtype=np.float64, copy=True)
    if (a.ndim != 2 or b.ndim != 2 or not 1 <= a.shape[0] <= 1000 or
            not 1 <= b.shape[0] <= 1000 or not 1 <= a.shape[1] <= 16 or a.shape[1] != b.shape[1]):
        raise ValueError('Expected bounded rows × matching features matrices.')
    if not np.isfinite(a).all() or not np.isfinite(b).all() or np.max(np.abs(a)) > 1e6 or np.max(np.abs(b)) > 1e6:
        raise ValueError('Inputs must be finite and bounded.')
    return a, b


def supplied_standardize(training, other):
    """Supplied inspection path, distinct from learner evidence."""
    a, b = validate_arrays(training, other)
    mean = a.mean(axis=0)
    spread = a.std(axis=0, ddof=0)
    scale = np.where(spread == 0, 1.0, spread)
    return (a - mean) / scale, (b - mean) / scale, mean, scale


def load_iris(path):
    with Path(path).open('rb') as stream:
        raw = stream.read(5978)
    if len(raw) != 5977 or hashlib.sha256(raw).hexdigest() != CSV_HASH:
        raise ValueError('Use the unchanged licensed 5,977-byte Iris CSV from the studio.')
    reader = csv.DictReader(io.StringIO(raw.decode('utf-8')))
    if reader.fieldnames != ['row_id', *FEATURES, 'species']:
        raise ValueError('Unexpected source columns.')
    records = list(reader)
    parts = {name: [] for name in ['training', 'development', 'evaluation']}
    seen, counts = set(), [0, 0, 0]
    for row in records:
        vector = tuple(float(row[key]) for key in FEATURES)
        if vector in seen:
            continue
        seen.add(vector)
        species = SPECIES.index(row['species'])
        bucket = counts[species] % 5
        counts[species] += 1
        part = 'evaluation' if bucket == 0 else 'development' if bucket == 1 else 'training'
        parts[part].append(dict(id=row['row_id'], x=vector, species=species))
    if [len(parts[k]) for k in parts] != [89, 30, 30]:
        raise ValueError('Unexpected grouped partition counts.')
    return parts


def make_report(parts, transform=supplied_standardize):
    train = np.array([row['x'] for row in parts['training']], dtype=np.float64)
    dev = np.array([row['x'] for row in parts['development']], dtype=np.float64)
    train_z, dev_z, mean, scale = transform(train.copy(), dev.copy())
    outputs = [train_z, dev_z, mean, scale]
    shapes = [train.shape, dev.shape, (4,), (4,)]
    if any(not isinstance(x, np.ndarray) or x.shape != shape or x.dtype != np.float64 or not np.isfinite(x).all()
           for x, shape in zip(outputs, shapes)) or (scale <= 0).any():
        raise ValueError('Transformation output violates shape, dtype or finite-value contract.')
    return dict(protocol=PROTOCOL, implementation='supplied' if transform is supplied_standardize else 'learner-unreviewed',
                source=dict(sha256=CSV_HASH, bytes=5977, source_rows=150, retained_groups=149,
                            removed_exact_measurement_duplicate_rows=1, units='cm',
                            attribution='Fisher, R. (1936). Iris. UCI Machine Learning Repository.',
                            doi='https://doi.org/10.24432/C56C76', license='CC BY 4.0'),
                split='Within each species, first-occurrence unique-vector index modulo 5: 0 evaluation, 1 development, others training. This studio protocol differs from the original two-way feasibility manifest.',
                partition_counts={name:len(rows) for name, rows in parts.items()},
                features=FEATURES, train_shape=list(train.shape), development_shape=list(dev.shape),
                dtype=str(train.dtype), mean_cm=mean.tolist(), scale_cm=scale.tolist(),
                constant_columns=(train.std(axis=0, ddof=0) == 0).tolist(),
                environment=dict(python=platform.python_version(), numpy=np.__version__, executable=sys.executable),
                training_standardized_column_mean=train_z.mean(axis=0).tolist(),
                development_rows=[dict(id=row['id'], species=SPECIES[row['species']], raw_cm=list(row['x']), z=z.tolist())
                                  for row,z in zip(parts['development'],dev_z)],
                limits='Descriptive public-data inspection only. No fitted classifier, no final score, no independent review. Supplied group definitions do not prove specimen identity.')


def save_artifacts(parts, report, destination):
    import matplotlib
    matplotlib.use('Agg')  # Save plots without a GUI or display server.
    import matplotlib.pyplot as plt
    destination = Path(destination)
    destination.mkdir(parents=False, exist_ok=False)  # Never silently replace a report.
    report['environment']['matplotlib'] = matplotlib.__version__
    (destination / 'report.json').write_text(json.dumps(report, indent=2, allow_nan=False) + '\n')
    with (destination / 'development.csv').open('w', newline='') as stream:
        writer = csv.writer(stream)
        writer.writerow(['id', 'species', *FEATURES, *['z_'+f for f in FEATURES]])
        for row in report['development_rows']:
            writer.writerow([row['id'], row['species'], *row['raw_cm'], *row['z']])
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.5), layout='constrained', sharex=True, sharey=True)
    colors, markers = ['#193e79', '#b42e39', '#796019'], ['o', 's', '^']
    for ax, partition in zip(axes, ['training', 'development']):
        for species, (color, marker) in enumerate(zip(colors, markers)):
            rows = [r['x'] for r in parts[partition] if r['species'] == species]
            ax.scatter([r[0] for r in rows], [r[2] for r in rows], label=SPECIES[species].replace('Iris-', ''),
                       color=color, marker=marker, s=28, alpha=.8)
        ax.set(xlim=(4,8.2), ylim=(0,7.2), xlabel='Sepal length (cm)', ylabel='Petal length (cm)',
               title=f'{partition.title()} · {len(parts[partition])} groups')
        ax.grid(alpha=.2)
    axes[1].legend(title='Species (shape and color)', loc='lower right', fontsize=8)
    fig.suptitle('Iris inspection · shared axes · no model predictions')
    fig.savefig(destination / 'inspection.png', dpi=150)
    plt.close(fig)
    (destination / 'README.txt').write_text('Plot compares two measured features, not model predictions. Both panels use the same axes. Species use shapes and colors. development.csv supplies numerical alternatives. Source: Fisher (1936), Iris, UCI, CC BY 4.0, https://doi.org/10.24432/C56C76. See report.json for exact bytes, split protocol and limitations.\n')


class ArrayPractice(unittest.TestCase):
    def test_known_columns_and_constant_feature(self):
        a = np.array([[1,10,5],[3,30,5]], dtype=np.float64)
        b = np.array([[5,50,5]], dtype=np.float64)
        az,bz,mean,scale = learner_standardize(a,b)
        np.testing.assert_allclose(mean,[2,20,5]); np.testing.assert_allclose(scale,[1,10,1])
        np.testing.assert_allclose(az,[[-1,-1,0],[1,1,0]]); np.testing.assert_allclose(bz,[[3,3,0]])
        np.testing.assert_array_equal(a,[[1,10,5],[3,30,5]])
        np.testing.assert_array_equal(b,[[5,50,5]])
    def test_development_does_not_fit(self):
        a=np.array([[1,2],[3,8],[5,14]], dtype=np.float64)
        first=learner_standardize(a,np.array([[10,20]],dtype=np.float64))
        second=learner_standardize(a,np.array([[100,200]],dtype=np.float64))
        for index in [0,2,3]: np.testing.assert_array_equal(first[index],second[index])
    def test_invalid_shapes_and_values(self):
        for a,b in [([],[[1]]),([[1]],[[1,2]]),([1,2],[[1]]),([[float('nan')]],[[1]]),([[1]],[[float('inf')]])]:
            with self.subTest(a=a,b=b), self.assertRaises(ValueError): learner_standardize(a,b)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--csv'); parser.add_argument('--out'); parser.add_argument('--check',action='store_true'); parser.add_argument('--learner',action='store_true')
    args=parser.parse_args()
    if args.check:
        result=unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(ArrayPractice))
        raise SystemExit(0 if result.wasSuccessful() else 1)
    if not args.csv or not args.out: parser.error('Provide --csv and a NEW --out directory, or --check.')
    parts=load_iris(args.csv)
    report=make_report(parts,learner_standardize if args.learner else supplied_standardize)
    save_artifacts(parts,report,args.out)
    print(f"Saved {args.out}. Implementation: {report['implementation']}. Human review pending.")

if __name__ == '__main__': main()
