# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import random
import numpy as np
from typing import Any, Dict, List

from glob import glob
from tqdm import tqdm 

from anydata.utils.types import is_list
from anydata.utils.read import read_json
from anydata.utils.data import make_list, flatten
from anydata.sync.sync_utils import get_seqs_subset


def merge_trees(trees: List['FolderTree']):
    """Merge multiple FolderTrees sequence-wise"""
    for tree in trees[1:]:
        trees[0].folder_tree.extend(tree.folder_tree)
    trees[0].prepare()
    return trees[0]


class FolderTree:
    """Creates a dataset tree folder structure for file loading"""
    def __init__(self, 
            path: Dict[str, Any],       # Dataset path
            context=(),                 # Length of temporal context
            stride=1,                   # Stride to use when loading sequences
            start=None,                 # Start from a given sequence index 
            finish=None,                # Finish at a given sequence index
            remove_files=None,          # Filter which files to remove
            keep_folders=None,          # Choose which sequences to keep
            remove_folders=None,        # Choose which sequences to remove
            folders_start=None,         # Keep sequences with a prefix filter
            folders_end=None,           # Keep sequences with a suffix filter
            folders_with=None,          # Keep folders with a certain string          
            folders_without=None,       # Keep folders without a certain string
            keep_files=None,            # Only keep certain filenames in each sequence 
            first_only=False,           # Only use the first sample in each sequence
            keep_first=None,            # Keep only the first few samples in each sequence
            pad_first=None,             # Pad the beginning of each sequence
            pad_last=None,              # Pad the end of each sequence
            crop_first=None,            # Crop the beginning of each scene 
            crop_last=None,             # Crop the end of each scene
            subset=None,                # Create a subset of available scenes
            single_slice=False,         # Treat each sequence as a single sample
            filter_camera=None,
        ):

        # Store context information
        self.context = list(context)
        if 0 not in self.context:
            self.context.append(0)
        self.num_context = 0 if len(self.context) == 0 else max(self.context) - min(self.context)
        self.with_context = self.num_context > 0
        self.min_context = 0 if not self.with_context else min(self.context)

        self.stride = stride
        self.single_slice = single_slice

        # Initialize empty folder tree
        self.folder_tree = []

        # Get folders from path dictionary
        folders = list(path.keys())
        # Choose which folders to keep
        if keep_folders is not None:
            folders = [f for f in folders if os.path.basename(f) in make_list(keep_folders)]
        # Choose which folders to remove
        if remove_folders is not None:
            folders = [f for f in folders if os.path.basename(f) not in make_list(remove_folders)]
        # Keep folders with a prefix filter
        if folders_start is not None:
            folders = flatten(
                [[f for f in folders if os.path.basename(f).startswith(str(start))] 
                for start in make_list(folders_start)])
        # Keep folders with a suffix filter
        if folders_end is not None:
            folders = flatten(
                [[f for f in folders if os.path.basename(f).endswith(str(start))] 
                for start in make_list(folders_end)])
        # Keep folders with a certain string
        if folders_with is not None:
            folders = flatten(
                [[f for f in folders if start in os.path.basename(f)] 
                for start in make_list(folders_with)])
        # Keep folders without a certain string
        if folders_without is not None:
            folders = flatten(
                [[f for f in folders if start not in os.path.basename(f)] 
                for start in make_list(folders_without)])

        # Populate folder tree
        for folder in folders:
            # Get and sort files in each folder
            files = path[folder] 
            # Start from a given sequence index
            if start is not None:
                files = files[start:]
            # Finish at a given sequence index
            if finish is not None:
                files = files[:finish]
            # Filter which files to remove
            if remove_files is not None:
                for remove in make_list(remove_files):
                    files = [f for f in files if remove not in f]
            # Filter which files to keep
            if keep_files is not None:
                for keep in make_list(keep_files):
                    files = [f for f in files if keep in f]
            # Apply stride
            if self.stride > 1:
                files = files[::self.stride]
            # Only store if there are more images than context
            if len(files) > self.num_context:
                self.folder_tree.append(files)

        # Only use the first sample in each sequence
        if first_only:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = self.folder_tree[i][:1]
        # Keep only the first few samples in each sequence
        if keep_first is not None:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = self.folder_tree[i][:keep_first]            
        # Pad the beginning of each sequence
        if pad_first is not None:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = [self.folder_tree[i][0]] * pad_first + self.folder_tree[i]
        # Pad the end of each sequence
        if pad_last is not None:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = self.folder_tree[i] + [self.folder_tree[i][-1]] * pad_last 
        # Crop the beginning of each scene
        if crop_first is not None:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = self.folder_tree[i][crop_first:]
        # Crop the end of each scene
        if crop_last is not None:
            for i in range(len(self.folder_tree)):
                self.folder_tree[i] = self.folder_tree[i][:-crop_last]

        # Create a subset of available scenes
        self.folder_tree: List[List[str]] = get_seqs_subset(self.folder_tree, subset)

        # Filter split to only contain valid sequences from a camera
        if filter_camera is not None:
            firsts = [f[0] for f in self.folder_tree]
            camera = firsts[0].split('/')[-2] # Get camera name
            metadatas = [read_json(
                os.path.dirname(os.path.dirname(os.path.dirname(f))) + '/metadata.json') 
                for f in firsts]
            filtered_folder_tree = []
            for i in range(len(self.folder_tree)):
                if camera in metadatas[i]['cameras']:
                    filtered_folder_tree.append(self.folder_tree[i])
            self.folder_tree = filtered_folder_tree

        # Prepare additional structures
        self.slices = self.total = None
        self.prepare()

    @staticmethod
    def from_list(trees: List['FolderTree']):
        """Merge FolderTrees from a list (sequence-wise)"""
        return merge_trees(trees)

    def __len__(self):
        """Dataset size"""
        if self.single_slice:
            return len(self.slices) - 1
        return self.total

    def prepare(self):
        """Prepare Folder Tree after sequences are defined"""
        # Get size of each folder
        self.slices = [len(folder) for folder in self.folder_tree]
        # Compensate for context size
        if self.with_context:
            self.slices = [s - self.num_context for s in self.slices]
        # Create cumulative size and get total
        self.slices = [0] + list(np.cumsum(self.slices))
        self.total = self.slices[-1]

    def get_idxs(self, idx):
        """Get folder and file indexes given dataset index"""
        if self.single_slice: return idx, 0 # If single slice, return the entire sequence starting from 0
        idx1 = np.searchsorted(self.slices, idx, side='right') - 1
        idx2 = idx - self.slices[idx1]
        return idx1, idx2

    def get_context_idxs(self, idx):
        """Get all context indexes"""
        idx1, idx2 = self.get_idxs(idx)
        return idx1, idx2, list(range(0, idx2)) + list(range(idx2 + 1, self.slices[idx1 + 1] - self.slices[idx1]))

    def get_all_idxs(self, idx):
        """Get target and context indexes"""
        idx1, idx2 = self.get_idxs(idx)
        return idx1, idx2, list(range(0, self.slice_length(idx1)))

    def get_proximity(self, idx1, offset):
        """Get folder and file indexes given dataset index"""
        return self.folder_tree[idx1][offset - self.min_context]

    def get_slice(self, idx):
        """Get dataset slice where the index resides"""
        idx1, idx2 = self.get_idxs(idx)
        return self.slices[idx1], self.slices[idx1+1]

    def get_item(self, idx, return_loc=False):
        """Return filename item given index"""
        idx1, idx2 = self.get_idxs(idx)
        item = {0: self.folder_tree[idx1][idx2 - self.min_context]}
        if return_loc:
            # print(idx2, self.min_context)
            return item, idx2 - self.min_context
        else:
            return item

    def slice_length(self, idx):
        """Get dataset slice length"""
        return self.slices[idx + 1]

    def get_context(self, idx, remove_target=True, context=None, interval=None):
        """Return forward context given index."""
        idx1, idx2 = self.get_idxs(idx)
        if context is None: context = self.context

        # If single slice, use entire context
        if self.single_slice: context = list(range(len(self.folder_tree[idx1])))

        if interval is not None:
            if isinstance(interval, list):
                curr = max(context)
                interv = interval[0] - curr
                interv = min(interv, len(self.folder_tree[idx1]))
                interv = interv // interval[1]
                interv = curr + int(random.choice(list(range(interv + 1)))) * interval[1]
                context = list(range(1, interv + 1)) + [0]
            else:
                context = list(range(1, interval + 1)) + [0]

        context = {ctx: self.folder_tree[idx1][idx2 - self.min_context + ctx] for ctx in context}
        if remove_target:
            for tgt in [0, (0, 0)]:
                if tgt in list(context.keys()):
                    context.pop(tgt)
        return context