import os
import pickle
from distutils.dir_util import copy_tree
from shutil import copyfile
from stack3d.study.components import StudyData, MouseData, TimepointData


def setup_study_collection_from_remote(path):

    """Scan the remote directory structure. Remote needs to be mounted first
    with Samba etc."""

    study_collection = []
    study_dirs = next(os.walk(path))[1]

    for study_dir in study_dirs:
        new_study = StudyData()
        new_study.raw_data_path = study_dir
        print "Set up study: ", study_dir
        mouse_dirs = next(os.walk(path + "/" + study_dir))[1]
        for mouse_dir in mouse_dirs:
            new_mouse = MouseData()
            new_mouse.raw_data_path = mouse_dir
            print "Set up mouse: ", mouse_dir

            cumulative_path = path + "/" + study_dir + "/" + mouse_dir
            if os.path.isdir(cumulative_path + "/Outputs_new_model"):
                new_mouse.extra_path = "/Outputs_new_model/"
            elif os.path.isdir(cumulative_path + "/Outputs"):
                new_mouse.extra_path = "/Outputs/"
            timepoint_dirs = next(os.walk(cumulative_path + new_mouse.extra_path))[1]
            for timepoint_dir in timepoint_dirs:
                if "TIF" not in timepoint_dir:
                    new_timepoint = TimepointData()
                    new_timepoint.raw_data_path = timepoint_dir
                    new_mouse.timepoints.append(new_timepoint)
            new_study.mice.append(new_mouse)
        study_collection.append(new_study)
    return study_collection


def setup_local_study_directory_structure(path, study_collection):

    for eachStudy in study_collection:
        study_dir = path + "/" + eachStudy.raw_data_path
        if not os.path.exists(study_dir):
            os.makedirs(study_dir)
        for eachMouse in eachStudy.mice:
            mouse_dir = study_dir + "/" + eachMouse.raw_data_path
            if not os.path.exists(mouse_dir):
                os.makedirs(mouse_dir)
            for eachTimePoint in eachMouse.timepoints:
                time_dir = mouse_dir + "/" + eachTimePoint.raw_data_path
                if not os.path.exists(time_dir):
                    os.makedirs(time_dir)


def copy_files(remote_path, local_path, study_collection):

    for eachStudy in study_collection:
        print "Copying Study: ", eachStudy.raw_data_path
        remote_study_dir = remote_path + "/" + eachStudy.raw_data_path
        local_study_dir = local_path + "/" + eachStudy.raw_data_path
        for eachMouse in eachStudy.mice:
            print "Copying Mouse: ", eachMouse.raw_data_path
            remote_mouse_dir = remote_study_dir + "/" + eachMouse.raw_data_path
            local_mouse_dir = local_study_dir + "/" + eachMouse.raw_data_path
            for eachTimePoint in eachMouse.timepoints:
                remote_time_dir = remote_mouse_dir + eachMouse.extra_path + "/" + eachTimePoint.raw_data_path
                local_time_dir = local_mouse_dir + "/" + eachTimePoint.raw_data_path

                if os.path.isdir(remote_time_dir + "/analysis_analysis"):
                    groups = ["csv", "histograms", "plots"]
                    for group in groups:
                        if os.path.isdir(remote_time_dir + "/analysis_analysis/" + group):
                            copy_tree(remote_time_dir + "/analysis_analysis/" + group,
                                      local_time_dir + "/analysis_analysis/" + group)

                    if os.path.isfile(remote_time_dir + "/analysis_analysis/skeleton.pkl"):
                        copyfile(remote_time_dir + "/analysis_analysis/skeleton.pkl", 
                                 local_time_dir + "/analysis_analysis/skeleton.pkl")
#             
work_dir = "/scratch/jgrogan/stack-working/study/"

samba_path = os.environ["XDG_RUNTIME_DIR"] + "/gvfs/"
raw_data_path = samba_path + "smb-share:server=imcore1.rob.ox.ac.uk,share=maths"
raw_data_path += "/Window Experiments/James/Network analysis/"

# study_collection = setup_study_collection_from_remote(raw_data_path)
# f = open(work_dir + "/study_collection.p", "wb")
# pickle.dump(study_collection, f)

f = open(work_dir + "/study_collection.p", 'r')
study_collection = pickle.load(f)

#setup_local_study_directory_structure(work_dir, study_collection)
copy_files(raw_data_path, work_dir, study_collection)