#!/usr/bin/env python

import climate
import numpy as np
import os

logging = climate.get_logger('amc2npy')

TAU = 2 * np.pi


def convert(filename, outdir=None):
    logging.info('converting %s', filename)
    frames = []
    with open(filename) as handle:
        frame = []
        for line in handle:
            line = line.strip()
            if line[0] in ':#':
                continue
            if line.isdigit():
                if frame:
                    frames.append(frame)
                frame = []
                assert int(line) == len(frames) + 1
                continue
            fields = line.split()
            segment = fields.pop(0)
            if segment == 'root':
                fields = fields[3:]
            frame.extend(map(float, fields))
        if frame:
            frames.append(frame)
    frames = TAU * np.array(frames) / 360
    output = filename.replace('.amc', '.npy')
    if outdir:
        output = os.path.join(outdir, os.path.basename(output))
    np.save(output, frames.astype('f'))
    logging.info('%s: wrote %s', output, frames.shape)


@climate.annotate(
    files='process these AMC files into NPY files',
    outdir='save output AMC files in this directory',
)
def main(*files, outdir='/data/cmu-classify/angle-data'):
    for f in files:
        try:
            convert(f, outdir)
        except:
            logging.exception('%s: error while processing', f)


if __name__ == '__main__':
    climate.call(main)
