from neo.io.proxyobjects import BaseProxy
from neo import Segment
import copy

def load_segment(segment, *args, **kwargs):
    """
    Utility function to load data of all child data object
    """

    new_segment = Segment(segment.name)

    def load_data(obj, *args, **kwargs):
        if isinstance(obj, BaseProxy):
            return obj.load(*args, **kwargs)
        else:
            return copy.deepcopy(obj)

    for st in segment.spiketrains:
        new_segment.spiketrains.append(load_data(st, load_waveforms=True))
    for event in segment.events:
        new_segment.events.append(load_data(event))
    for epoch in segment.epochs:
        new_segment.epochs.append(load_data(epoch))
    for anasig in segment.analogsignals:
        new_segment.analogsignals.append(load_data(anasig))
    for irrsig in segment.irregularlysampledsignals:
        new_segment.irregularlysampledsignals.append(load_data(irrsig))

    for c in new_segment.children:
        c.segment = new_segment

    return new_segment