#!/usr/bin/env python3


def generate_collection_class(obj_name, obj_attrs):
    src = []

    src += f'''\
struct {obj_name};

class {obj_name}Collection {{
  public:
    class iter {{
      public:
        iter({obj_name}Collection* collection, size_t idx)
          :collection(collection), idx(idx) {{ }}
        iter operator++() {{ ++idx; return *this; }}
        bool operator!=(const iter & other) {{ return idx != other.idx; }}
        {obj_name} operator*();
      private:
        {obj_name}Collection* collection;
        size_t idx;
    }};
    TrackingDataSet* tds;

'''.splitlines()

    for field in obj_attrs['fields']:
        name = field['name']
        type_ = field['type']
        src.append(f'    Value<vector<{type_}>>* val_{name};')
        src.append(f'    bool {name}_loaded;')

    src.append(f'\n    {obj_name}Collection() {{ }}\n')

    src.append('    void init(TrackingDataSet* tds){')
    src.append('        this->tds = tds;')
    src.append('    }\n')
    first_obj_name = list(obj_attrs['fields'])[0]['name']
    first_obj_type = list(obj_attrs['fields'])[0]['type']
    prefix = obj_attrs['treename_prefix']
    src.append(f'''\
    size_t size() {{
        if (!this->{first_obj_name}_loaded) {{
            this->val_{first_obj_name} = this->tds->track_branch_obj<vector<{first_obj_type}>>("{prefix}_{first_obj_name}");
            this->{first_obj_name}_loaded = true;
        }}
        return (*this->val_{first_obj_name})().size();
    }}
\n''')
    src.append(f'    {obj_name} operator[](size_t);')
    src.append('    iter begin() { return iter(this, 0); }')
    src.append('    iter end() { return iter(this, size()); }')
    src.append('};')

    src += f'''
struct {obj_name} {{
    {obj_name}Collection* collection;
    size_t idx;
    {obj_name}({obj_name}Collection* collection, const size_t idx)
      :collection(collection), idx(idx) {{ }}\n
'''.splitlines()

    for field in obj_attrs['fields']:
        name = field['name']
        type_ = field['type']
        # Because vector<bool> is packed, a temporary object is created so we can't return a reference.
        ret_type = f'{type_}&' if type_ != 'bool' else type_
        src.append(f'''\
    const {ret_type} {name}() const {{
        if (!collection->{name}_loaded) {{
            collection->val_{name} = collection->tds->track_branch_obj<vector<{type_}>>("{prefix}_{name}");
            collection->{name}_loaded = true;
        }}
        return (*collection->val_{name})().at(idx);
    }}
''')
    src.append('};')

    src.append(f'''
bool operator==(const {obj_name}& obj1, const {obj_name}& obj2) {{
    return obj1.idx == obj2.idx;
}}

{obj_name} {obj_name}Collection::iter::operator*() {{
    return {{collection, idx}};
}}
{obj_name} {obj_name}Collection::operator[](size_t idx) {{
    return {{this, idx}};
}}
''')
    return '\n'.join(src)


def generate_header(input_filename, output_filename):
    from datetime import datetime

    return f'''\
/** {output_filename} created on {datetime.now()} by generate_class.py
 * AVOID EDITING THIS FILE BY HAND!! Instead edit {input_filename} and re-run
 * generate_class.py
 */
#include "filval.hpp"
#include "root_filval.hpp"

#include<cmath>

#include "TrackingNtuple.h"

using namespace std;
using namespace fv;
using namespace fv_root;

typedef TreeDataSet<TrackingNtuple> TrackingDataSet;
'''


if __name__ == '__main__':
    import argparse
    import yaml
    parser = argparse.ArgumentParser()
    add = parser.add_argument
    add('input_file', help='An input YAML file defining the objects to generate')

    args = parser.parse_args()
    classes = []
    with open(args.input_file) as fi:
        for obj, attrs in yaml.load(fi).items():
            classes.append(generate_collection_class(obj, attrs))
    output_filename = args.input_file.replace('.yaml', '.hpp')
    with open(output_filename, 'w') as fo:
        fo.write(generate_header(args.input_file, output_filename))
        for class_ in classes:
            fo.write(class_)