import pyrtl

pyrtl.reset_working_block()

class SimplePipeline(object):
    def __init__(self):
        self._pipeline_register_map = {}
        self._current_stage_num = 0
        stage_list = [method for method in dir(self) if method.startswith('stage')]
        for stage in sorted(stage_list):
            stage_method = getattr(self, stage)
            stage_method()
            self._current_stage_num += 1
    
    def __getattr__(self, name):
        try:
            return self._pipeline_register_map[self._current_stage_num][name]
        except KeyError:
            raise pyrtl.PyrtlError(
                'error, no pipeline register "%s" defined for stage %d'
                % (name, self._current_stage_num))

    def __setattr__(self, name, value):
        if name.startswith('_'):
            # do not do anything tricky with variables starting with '_'
            object.__setattr__(self, name, value)
        else:
            next_stage = self._current_stage_num + 1
            pipereg_id = str(self._current_stage_num) + 'to' + str(next_stage)
            rname = 'pipereg_' + pipereg_id + '_' + name
            new_pipereg = pyrtl.Register(bitwidth=len(value), name=rname)
            if next_stage not in self._pipeline_register_map:
                self._pipeline_register_map[next_stage] = {}
            self._pipeline_register_map[next_stage][name] = new_pipereg
            new_pipereg.next <<= value    

"""
Explanation: Example 5: Making use of PyRTL and Introspection.
The following example shows how pyrtl can be used to make some interesting
hardware structures using python introspection.  In particular, this example
makes a N-stage pipeline structure.  Any specific pipeline is then a derived
class of SimplePipeline where methods with names starting with "stage" are
stages, and new members with names not starting with "_" are to be registered
for the next stage.
Pipeline builder with auto generation of pipeline registers.
End of explanation
"""


class SimplePipelineExample(SimplePipeline):
    def __init__(self):
        self._loopback = pyrtl.WireVector(1, 'loopback')
        super(SimplePipelineExample, self).__init__()

    def stage0(self):
        self.n = ~ self._loopback

    def stage1(self):
        self.n = self.n

    def stage2(self):
        self.n = self.n

    def stage3(self):
        self.n = self.n

    def stage4(self):
        self._loopback <<= self.n

"""
Explanation: A very simple pipeline to show how registers are inferred.
End of explanation
"""


simplepipeline = SimplePipelineExample()

sim_trace = pyrtl.SimulationTrace()
sim = pyrtl.Simulation(tracer=sim_trace)

for cycle in range(15):
    sim.step({})
sim_trace.render_trace()

"""
Explanation: Simulation of the core
End of explanation
"""


