from caffe_all import *

def parseProtoString(s):
    from google.protobuf import text_format
    proto_net = pb.NetParameter()
    text_format.Merge(s, proto_net)
    return proto_net


def get_param(l, exclude=set(['top', 'bottom', 'name', 'type'])):
    if not hasattr(l,'ListFields'):
        if hasattr(l,'__delitem__'):
            return [get_param(i) for i in l]
        return l
    r = dict()
    for f, v in l.ListFields():
        if f.name not in exclude:
            r[f.name] = get_param(v, [])
    return r

class ProtoDesc:
    def __init__(self, prototxt):
        from os import path
        self.prototxt = prototxt
        self.parsed_proto = parseProtoString(open(self.prototxt, 'r').read())
        # Guess the input dimension
        self.input_dim = (3, 227, 227)
        net = self.parsed_proto
        if len(net.input_dim) > 0:
            self.input_dim = net.input_dim[1:]
        else:
            lrs = net.layer
            cs = [l.transform_param.crop_size for l in lrs
                  if l.HasField('transform_param')]
            if len(cs):
                self.input_dim = (3, cs[0], cs[0])

    def __call__(self, clip=None, **inputs):
        from collections import OrderedDict
        net = self.parsed_proto
        blobs = OrderedDict(inputs)
        for l in net.layer:
            if l.type not in ['Data', 'ImageData']:
                in_place = l.top == l.bottom
                param = get_param(l)
                tops = getattr(L, l.type)(*[blobs[b] for b in l.bottom],
                                          ntop=len(l.top), in_place=in_place,
                                          name=l.name,
                                          **param)
                if len(l.top) <= 1:
                    tops = [tops]
                for i, t in enumerate(l.top):
                    blobs[t] = tops[i]
            if l.name == clip:
                break
        return list(blobs.values())[-1]
