# Copyright, 1999, Regents of the University of California
# Please see file Legal.htm
import types, string

class FortranType:
    "Description of a Fortran type."
    def __init__ (self, name, kind=None):
        assert type(name) == types.StringType
        assert kind is None or type(kind) == types.StringType
        self.name = name
        self.kind = kind

    def __repr__ (self):
        if self.name == 'character' and self.kind:
            s = 'character*(' + self.kind + ')'
        else:
            s = self.name
            if self.kind:
                s = s + '(' + self.kind + ')'
        return s

class FortranAttributes:
    "Attributes a Fortran entity can have"
    def __init__(self, type, intent, allocatable):
        assert is_FortranType(type)
        assert intent == 'in' or intent == 'out' or intent == 'inout' \
                or intent == 'temporary' or intent == 'valued'
        self.intent = intent
        self.type = type
        self.allocatable = allocatable

    def __repr__(self):
        s = repr(self.type) + ", intent = " + self.intent
        if self.allocatable: s = s + ", allocatable"
        return s

class FortranProcedure:
    "Interface info to a Fortran subroutine or function"
    def __init__(self, head, declarations):
        self.head = head
        self.dict = {}
        self.dict[self.routine_name()] = None
        for a in self.arguments():
            self.dict[a] = None
        for d in declarations:
            self.dict[d.name] = d
        for name in self.dict.keys():
            if name == self.routine_name() and not self.is_function():
                if not (self.dict[name] is None):
                    print name + " is declared in the body but it is a subroutine."
                    raise SystemExit, 1
                self.dict[name] = void_declaration(name)
            else:
                if self.dict[name] is None:
                    print name + " is not given a type in routine " + self.routine_name()
                    raise SystemExit, 1
        if self.is_function():
            self.dict[self.routine_name()].set_intent("out")

    def __repr__(self):
        s = repr(self.head) + "\n" 
        for x in self.declarations():
            s = s + "     " + repr(x) + "\n"
        return s

    def routine_name (self):
        "name of the routine"
        return self.head[1]

    def arguments (self):
        "names of the arguments"
        return self.head[2]

    def python_inputs(self):
        "list of arguments whose intent is in or inout"
        result = []
        for x in self.arguments():
            k = self.dict[x].intent
            if k == 'in' or k == 'inout':
                result.append(x)
        return result

    def valued (self):
        "list of arguments whose intent is valued"
        result = []
        for x in self.arguments():
            k = self.dict[x].intent
            if k == 'valued':
                result.append(x)
        return result

    def temporaries(self):
        "list of arguments whose intent is temporary"
        result = []
        for x in self.arguments():
            k = self.dict[x].intent
            if k == "temporary":
                result.append(x)
        return result

    def in_and_outs(self):
        "list of arguments whose intent is inout"
        result = []
        for x in self.arguments():
            k = self.dict[x].intent
            if k == "inout":
                result.append(x)
        return result

    def python_outputs(self):
        """list of arguments of type out, preceded by name of function
           if it is a function"""
# note that inouts are not outputs
        result = []
        if self.is_function():
            result.append(self.routine_name())
        for x in self.arguments():
            k = self.dict[x].intent
            if k == 'out':
                result.append(x)
        return result

    def routine_type (self):
        "subroutine or function?"
        return self.head[0]

    def return_type (self):
        "type returned by this function"
        if self.is_function():
            return self.declaration(self.routine_name()).type
        else:
            return FortranType("none")

    def is_function (self):
        "is this a function?"
        return self.routine_type() == 'function'

    def head_comments (self, indent=0):
        "Comments connected with the procedure itself."
        cs = []
        sep = ' '*indent
        for c in self.head[3]:
            cc = c
            if cc[0:1] == '!': cc = cc[1:]
            if cc[0:1] == ' ': cc = cc[1:]
            cs.append(sep + cc)
        return sep + string.join(cs, '\n')

    def declaration(self, name):
        "the declaration object for the entity with the given name"
        return self.dict[name]

    def declarations(self):
        "The list of declarations"
        result = []
        if self.is_function():
            result.append(self.declaration(self.routine_name()))
        for s in self.arguments():
            result.append(self.declaration(s))
        return result

    def make_procedure(self):
        routine_type = 'procedure'
        routine_name = self.head[1]
        argument_name = routine_name + 'arg'
        arguments = [argument_name] + self.head[2]
        declaration = FortranDeclaration (argument_name,[])
        declaration.set_intent ("out")
        declaration.set_type (self.dict[routine_name].type)
        self.dict[argument_name] = declaration
        self.head = routine_type, self.head[1], arguments, self.head[3]
        type = FortranType ("none")
        self.dict[routine_name].set_type (type)

    def add_argument(self, name, type, intent, preceding):
        arguments = self.head[2]
        i = arguments.index(preceding)
        arguments = arguments[:i+1] + [name] + arguments[i+1:]
        self.head = self.head[0], self.head[1], arguments, self.head[3]
        declaration = FortranDeclaration (name,[])
        declaration.set_intent (intent)
        declaration.set_type (type)
        self.dict[name] = declaration

unknown_type = FortranType('unknown-type')

class FortranDeclaration:
    "Declaration of a variable"
    def __init__(self, name, dimlist):
        assert type(name) == types.StringType
        assert type(dimlist) == types.ListType
        self.name = name
        self.dimlist = dimlist
        self.type = unknown_type
        self.intent = 'unknown-intent'
        self.allocatable = 0
        self.comment = ''

    def set_info(self, attributes):
        assert is_FortranAttributes(attributes)
        self.set_type(attributes.type)
        self.set_intent(attributes.intent)
        self.set_allocatable(attributes.allocatable)

    def set_comment(self, comment):
        assert type(comment) == types.StringType
        self.comment = comment

    def set_type(self, type):
        self.type = type

    def set_intent (self, intent):
        self.intent = intent

    def set_allocatable (self, allocatable):
        self.allocatable = allocatable

    def dimstr(self):
        "String repr of dimlist "
        if self.dimlist:
            s = string.join (self.dimlist, ', ')
        else:
            s = ''
        return s

    def rank (self):
        "Number of dimensions."
        return len(self.dimlist)

    def dimension (self, i):
        "i'th dimension, 0 based."
        assert i <= self.rank()
        assert i >= 0
        return self.dimlist[i]

    def __repr__(self):
        s = repr(self.type) + ', intent(' + self.intent + '):: ' + self.name 
        if self.dimlist:
            s = s + ' (' + self.dimstr() + ')'
        if self.allocatable:
            s = s + ' allocatable'
        return s

def void_declaration (name):
    d = FortranDeclaration (name, [])
    d.set_intent("out")
    d.set_type(FortranType("none"))
    return d

def is_FortranAttributes (s):
    return type(s) == types.InstanceType and s.__class__ == FortranAttributes

def is_FortranProcedure (s):
    return type(s) == types.InstanceType and s.__class__ == FortranProcedure

def is_FortranType (s):
    return type(s) == types.InstanceType and s.__class__ == FortranType

def is_FortranDeclaration (s):
    return type(s) == types.InstanceType and s.__class__ == FortranDeclaration