from typing import Union
from typing import List
import pickle


def parse_list(data: Union[str, List],
               indicator: int = 0,
               query=['MCF7']) -> List:
  """Filter the data based on compound, cell line, dose or time

  This function takes the directory of dataset, indicator that indicates
  whether you want to subset the data based on cell line, compound, dose, or time
  and a list which shows what part of the data you want to keep.
  The output will be a list of desired parsed dataset.


  Parameters
  ----------
  data: Union[str, List]
    the data can be a string which is the directory of the dataset.
    dataset should be a pickle file. e.g., valid argument is something like this:
    './Data/level3_trt_cp_landmark.pkl'
    or it can be a list which contains the gene expression and metadata.
    It must be a list of tuples with the following format:
    line[0]:(cell_line, drug, drug_type, does, does_type, time, time_type)
    line[1]: 978 or 12328-dimensional Vector(Gene_expression_profile)

  indicator: int
    it must be an integer from 0 1 2 and 3 that shows whether
    we want to retrieve the data based on cells, compound or dose.
    0: cell_lines
    1:compounds
    2:doses
    3:time
    Default=0 (cell_lines)

  query: List
    list of cells or compounds or doses that we want to retrieve.
    The list depends on the indicator. If the indicator is 0, you should enter the
    list of desired cell lines and so on. Default=['MCF7']

  Returns
  -------
  parse_data: List
    A list containing data that belongs to desired list.
  """

  assert isinstance(indicator, int), "The indicator must be an int object"
  assert indicator in [0, 1, 2,
                       3], "You should choose indicator from 0, 1, 2 range"
  assert isinstance(query, list), "The parameter query must be a list"

  print("=================================================================")
  print("Data Loading..")

  assert isinstance(data,
                    (str, list)), "The data should be string or list object"
  if isinstance(data, str):
    with open(data, "rb") as f:
      train = pickle.load(f)
  else:
    assert isinstance(data, list), "The data must be a list object"
    train = data

  mapping = {0: 0, 1: 1, 2: 3, 3: 5}
  k = mapping[indicator]
  mapping_name = {0: 'cell_lines', 1: 'compounds', 2: 'doses', 3: 'time'}

  print("Number of Train Data: {}".format(len(train)))
  print("You are parsing the data base on {}".format(mapping_name[indicator]))

  parse_data = [line for line in train if line[0][k] in query]

  print("Number of Data after parsing: {}".format(len(parse_data)))
  return parse_data