processing.py

"""
Function that ignite from main_window, Dialogs to read data files for data,
channels, datatype
"""

import os
import json
import re
import traceback
from pathlib import Path
from typing import List, Set, Optional, Dict, Tuple

from PySide2.QtCore import QEventLoop, Qt
from PySide2.QtGui import QCursor
from PySide2.QtWidgets import QTextBrowser, QApplication
from obspy.core import read as read_ms
from obspy.io.reftek.core import Reftek130Exception

from sohstationviewer.database.extract_data import get_signature_channels
from sohstationviewer.model.data_type_model import DataTypeModel
from sohstationviewer.model.handling_data import (
    read_mseed_chanids_from_headers)

from sohstationviewer.controller.util import (
    validate_file, display_tracking_info
)

from sohstationviewer.view.util.enums import LogType


def load_data(data_type: str, tracking_box: QTextBrowser, dir_list: List[str],
              list_of_rt130_paths: List[Path],
              req_wf_chans: List[str] = [], req_soh_chans: List[str] = [],
              read_start: Optional[float] = None,
              read_end: Optional[float] = None) -> DataTypeModel:
    """
    Load the data stored in list_of_dir and store it in a DataTypeModel object.
    The concrete class of the data object is based on dataType. Run on the same
    thread as its caller, and so will block the GUI if called on the main
    thread. It is advisable to use model.data_loader.DataLoader to load data
    unless it is necessary to load data in the main thread (e.g. if there is
    a need to access the call stack).

    :param data_type: type of data read
    :param tracking_box: widget to display tracking info
    :param dir_list: list of directories selected by users
    :param list_of_rt130_paths: list of rt130 directories selected by users
    :param req_wf_chans: requested waveform channel list
    :param req_soh_chans: requested soh channel list
    :param read_start: start time of read data
    :param read_end: finish time of read data
    :return data_object: object that keep the data read from
        list_of_dir
    """
    data_object = None
    if list_of_rt130_paths == []:
        for d in dir_list:
            if data_object is None:
                try:
                    data_object = DataTypeModel.create_data_object(
                        data_type, tracking_box, d, [],
                        req_wf_chans=req_wf_chans, req_soh_chans=req_soh_chans,
                        read_start=read_start, read_end=read_end)
                except Exception:
                    fmt = traceback.format_exc()
                    msg = f"Dir {d} can't be read due to error: {str(fmt)}"
                    display_tracking_info(tracking_box, msg, LogType.WARNING)

                # if data_object.has_data():
                #     continue
                # If no data can be read from the first dir, throw exception
                # raise Exception("No data can be read from ", d)
            # TODO: will work with select more than one dir later
            # else:
            #     data_object.readDir(d)

    else:
        try:
            data_object = DataTypeModel.create_data_object(
                data_type, tracking_box, [''], list_of_rt130_paths,
                req_wf_chans=req_wf_chans, req_soh_chans=req_soh_chans,
                read_start=read_start, read_end=read_end)
        except Exception:
            fmt = traceback.format_exc()
            msg = f"RT130 selected can't be read due to error: {str(fmt)}"
            display_tracking_info(tracking_box, msg, LogType.WARNING)

    if data_object is None:
        msg = "No data object created. Check with implementer"
        display_tracking_info(tracking_box, msg, LogType.WARNING)
    return data_object


def read_mseed_channels(tracking_box: QTextBrowser, list_of_dir: List[str],
                        on_unittest: bool = False
                        ) -> Set[str]:
    """
    Scan available for SOH channels (to be used in channel preferences dialog).
        Since channels for RT130 is hard code, this function won't be applied
        for it.
    Note that Mass position channels are excluded because the default
        include_mp123zne and include_mp456uvw for MSeed are False
    :param tracking_box: widget to display tracking info
    :param list_of_dir: list of directories selected by users
    :param on_unittest: flag to avoid cursor code to display tracking_info
    :return data_object.channels: set of channels present in listofDir
    """
    soh_chan_ids = set()
    mass_pos_chan_ids = set()
    wf_chan_ids = set()
    spr_gr_1_chan_ids = set()
    count = 0
    if not on_unittest:
        QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
    display_tracking_info(
        tracking_box, "Start reading channel's header", LogType.INFO)
    for d in list_of_dir:
        for path, sub_dirs, files in os.walk(d):
            if 'System Volume Information' in path:
                continue
            for file_name in files:
                if not on_unittest:
                    QApplication.processEvents(
                        QEventLoop.ExcludeUserInputEvents
                    )
                path2file = Path(path).joinpath(file_name)
                if not validate_file(path2file, file_name):
                    continue
                count += 1
                if count % 10 == 0:
                    display_tracking_info(
                        tracking_box, f'Read {count} file headers/ SOH files',
                        LogType.INFO)

                ret = read_mseed_chanids_from_headers(path2file, file_name)
                soh_chan_ids.update(ret[0])
                mass_pos_chan_ids.update(ret[1])
                wf_chan_ids.update(ret[2])
                spr_gr_1_chan_ids.update(ret[3])
    if not on_unittest:
        QApplication.restoreOverrideCursor()
    return sorted(list(soh_chan_ids)), sorted(list(mass_pos_chan_ids)), \
        sorted(list(wf_chan_ids)), sorted(list(spr_gr_1_chan_ids))


def detect_data_type(list_of_dir: List[str]) -> Optional[str]:
    """
    Detect data type for the given directories using get_data_type_from_file
    :param list_of_dir: list of directories selected by users
    :return:
        + if there are more than one data types detected,
            return None with a warning message
        + if only Unknown data type detected,
            return None with a warning message
        + if data type found, return data_type,
    """
    sign_chan_data_type_dict = get_signature_channels()

    dir_data_type_dict = {}
    for d in list_of_dir:
        data_type = "Unknown"
        for path, subdirs, files in os.walk(d):
            for file_name in files:
                path2file = Path(path).joinpath(file_name)
                if not validate_file(path2file, file_name):
                    continue
                ret = get_data_type_from_file(path2file,
                                              sign_chan_data_type_dict)
                if ret is not None:
                    data_type, chan = ret
                    break
            if data_type != "Unknown":
                break
        if data_type == "Unknown":
            dir_data_type_dict[d] = ("Unknown", '_')
        else:
            dir_data_type_dict[d] = (data_type, chan)
    data_type_list = {d[0] for d in dir_data_type_dict.values()}
    if len(data_type_list) > 1:
        dir_data_type_str = json.dumps(dir_data_type_dict)
        dir_data_type_str = re.sub(r'\{|\}|"', '', dir_data_type_str)
        dir_data_type_str = re.sub(r'], ', ']\n', dir_data_type_str)
        msg = (f"There are more than one types of data detected:\n"
               f"{dir_data_type_str}\n\n"
               f"Please have only data that related to each other.")
        raise Exception(msg)

    elif data_type_list == {'Unknown'}:
        msg = ("There are no known data detected.\n"
               "Please select different folder(s).")
        raise Exception(msg)

    return list(dir_data_type_dict.values())[0][0]


def get_data_type_from_file(
        path2file: Path,
        sign_chan_data_type_dict: Dict[str, str]
) -> Optional[Tuple[str, str]]:
    """
    + Try to read mseed data from given file
        if catch TypeError: no data type detected => return None
        if catch Reftek130Exception: data type => return data type RT130
        otherwise data type is mseed which includes: q330, pegasus, centaur
    + Continue to identify data type for a file by checking if the channel
    in that file is a unique channel of a data type.
    :param path2file: absolute path to processed file
    :param sign_chan_data_type_dict: dict of unique chan for data
        type
    :return: detected data type, channel from which data type is detected
    """
    try:
        stream = read_ms(path2file)
    except TypeError:
        return
    except Reftek130Exception:
        return 'RT130', '_'

    for trace in stream:
        chan = trace.stats['channel']
        if chan in sign_chan_data_type_dict.keys():
            return sign_chan_data_type_dict[chan], chan