#!/usr/bin/env python3
import yaml
import numpy as np
import sys


def features(x, y):
    return np.array([1.0, x, y, x * x, x * y, y * y], dtype=np.float64)


def main():
    sample_path = sys.argv[1] if len(sys.argv) > 1 else 'samples.yaml'
    model_path = sys.argv[2] if len(sys.argv) > 2 else 'model.yaml'

    with open(sample_path, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)

    samples = data['samples']
    if len(samples) < 6:
        print('样本太少，至少需要 6 个，建议 9~16 个')
        return

    X = []
    Y = []

    for s in samples:
        X.append(features(s['x_cam'], s['y_cam']))
        Y.append(s['arm_cmd'])

    X = np.array(X, dtype=np.float64)            # [N, 6]
    Y = np.array(Y, dtype=np.float64)            # [N, cmd_dim]

    # 最小二乘
    W, _, _, _ = np.linalg.lstsq(X, Y, rcond=None)   # [6, cmd_dim]
    Y_pred = X @ W
    err = Y_pred - Y
    mae = np.mean(np.abs(err), axis=0)
    rmse = np.sqrt(np.mean(err ** 2, axis=0))

    model = {
        'feature_order': ['1', 'x', 'y', 'x2', 'xy', 'y2'],
        'weights': W.tolist(),
        'cmd_dim': int(Y.shape[1]),
        'sample_count': int(len(samples)),
        'mae_per_dim': mae.tolist(),
        'rmse_per_dim': rmse.tolist(),
        'z_ref_median': float(np.median([s['z_cam'] for s in samples]))
    }

    with open(model_path, 'w', encoding='utf-8') as f:
        yaml.safe_dump(model, f, allow_unicode=True, sort_keys=False)

    print(f'模型已保存到: {model_path}')
    print('每维 MAE:', mae)
    print('每维 RMSE:', rmse)


if __name__ == '__main__':
    main()