path: root/r_qtl/r_qtl2.py

"""The R/qtl2 parsing and processing code."""
import io
import csv
import json
from pathlib import Path
from zipfile import ZipFile
from functools import reduce, partial
from typing import Union, Iterator, Iterable, Callable, Optional

import yaml

from functional_tools import take, chain

from r_qtl.errors import InvalidFormat, MissingFileError

FILE_TYPES = (
    "geno", "founder_geno", "pheno", "covar", "phenocovar", "gmap", "pmap",
    "phenose")

def control_data(zfile: ZipFile) -> dict:
    """Retrieve the control file from the zip file info."""
    files = tuple(filename
                  for filename in zfile.namelist()
                  if (filename.endswith(".yaml") or filename.endswith(".json")))
    num_files = len(files)
    if num_files == 0:
        raise InvalidFormat("Expected a json or yaml control file.")

    if num_files > 1:
        raise InvalidFormat("Found more than one possible control file.")

    return {
        "na.strings": ["NA"],
        "comment.char": "#",
        "sep": ",",
        **{
            f"{key}_transposed": False for key in FILE_TYPES
        },
        **(json.loads(zfile.read(files[0]))
            if files[0].endswith(".json")
            else yaml.safe_load(zfile.read(files[0])))
    }

def replace_na_strings(cdata, val):
    """Replace values indicated in `na.strings` with `None`."""
    return (None if val in cdata.get("na.strings", ["NA"]) else val)

def with_non_transposed(zfile: ZipFile,
                        member_key: str,
                        cdata: dict,
                        process_value: Callable[
                            [dict], dict] = lambda val: val) -> Iterator[dict]:
    """Process non-transposed file values

    Arguments:
    zfile: A zipfile object from opening a R/qtl2 bundle.
    member_key: A key to retrieve the member file to process from the file.
    cdata: The control data from the R/qtl2 bundle read from the JSON/YAML file.
    process_value: A function to process the values from the file.
    """
    def not_comment_line(line):
        return not line.startswith(cdata.get("comment.char", "#"))

    sep = cdata.get("sep", ",")
    with zfile.open(cdata[member_key]) as innerfile:
        try:
            wrapped_file = io.TextIOWrapper(innerfile)
            firstrow = tuple(
                field.strip() for field in
                next(filter(not_comment_line, wrapped_file)).strip().split(sep))
            id_key = firstrow[0]
            wrapped_file.seek(0)
            reader = csv.DictReader(filter(not_comment_line, wrapped_file),
                                    delimiter=sep)
            for row in reader:
                processed = process_value(row)
                yield {
                    "id": processed[id_key],
                    **{
                        key: value
                        for key, value in processed.items()
                        if key != id_key
                    }
                }
        except StopIteration as exc:
            raise InvalidFormat("The file has no rows!") from exc

def __make_organise_by_id__(id_key):
    """Return a function to use with `reduce` to organise values by some
    identifier."""
    def __organiser__(acc, item):
        row = acc.get(item[id_key], {})
        return {**acc, item[id_key]: {**row, **item}}
    return __organiser__

def __batch_of_n__(iterable: Iterable, num):
    """Return a batch of `num` items or less from the `iterable`."""
    while True:
        items = take(iterable, num)
        if len(items) <= 0:
            break
        yield items

def with_transposed(zfile: ZipFile,
                    member_key: str,
                    cdata: dict,
                    process_value: Callable[
                        [str, tuple[str, ...], tuple[str, ...]],
                        tuple[dict, ...]]) -> Iterator[dict]:
    """Process transposed file values

    Arguments:
    zfile: A zipfile object from opening a R/qtl2 bundle.
    member_key: A key to retrieve the member file to process from the file.
    cdata: The control data from the R/qtl2 bundle read from the JSON/YAML file.
    process_value: A function to process the values from the file.
    """
    with zfile.open(cdata[member_key]) as innerfile:
        lines = (tuple(field.strip() for field in
                       line.strip().split(cdata.get("sep", ",")))
                 for line in
                 filter(lambda line: not line.startswith("#"),
                        io.TextIOWrapper(innerfile)))
        try:
            id_line = next(lines)
            id_key, headers = id_line[0], id_line[1:]
            for _key, row in reduce(# type: ignore[var-annotated]
                    __make_organise_by_id__(id_key),
                    (row
                     for batch in __batch_of_n__(lines, 300)
                     for line in batch
                     for row in process_value(id_key, headers, line)),
                    {}).items():
                yield {
                    "id": row[id_key],
                    **{
                        key: value
                        for key, value in row.items()
                        if key != id_key
                    }}
        except StopIteration:
            pass

def make_process_data_geno(cdata) -> tuple[
        Callable[[dict], dict],
        Callable[[str, tuple[str, ...], tuple[str, ...]],
                 tuple[dict, ...]]]:
    """Build functions to process genotype data."""
    def replace_genotype_codes(val):#pylint: disable=[redefined-outer-name]
        # The rewrite will probably make this obsolete.
        return cdata["genotypes"].get(val, val)

    def __non_transposed__(row: dict) -> dict:
        return {
            key: chain(value, replace_genotype_codes,
                       partial(replace_na_strings, cdata))
            for key,value in row.items()
        }
    def __transposed__(id_key: str,
                       ids: tuple[str, ...],
                       vals: tuple[str, ...]) -> tuple[dict, ...]:
        return tuple(
            dict(zip(
                [id_key, vals[0]],
                (chain(item, replace_genotype_codes, partial(replace_na_strings, cdata))
                 for item in items)))
            for items in zip(ids, vals[1:]))
    return (__non_transposed__, __transposed__)

def replace_sex_info(val, cdata: dict):
    """Replace sex information in files with values in the control data."""
    sex_info = cdata.get("sex", False)
    if bool(sex_info):
        return sex_info.get(val, val)
    return val

def replace_cross_info(val, cdata: dict):
    """
    Replace cross information in files with the values in the control data.
    """
    cross_info = cdata.get("cross_info", False)
    if bool(cross_info):
        return cross_info.get(val, val)
    return val

def make_process_data_covar(cdata) -> tuple[
        Callable[[dict], dict],
        Callable[[str, tuple[str, ...], tuple[str, ...]],
                 tuple[dict, ...]]]:
    """Build functions to process sex and cross information in covar files."""
    rep_sex_info = partial(replace_sex_info, cdata=cdata)
    rep_cross_info = partial(replace_cross_info, cdata=cdata)
    def non_transposed(row: dict) -> dict:
        return {
            key: chain(value, rep_sex_info, rep_cross_info)
            for key,value in row.items()
        }
    def transposed(id_key: str,
                   ids: tuple[str, ...],
                   vals: tuple[str, ...]) -> tuple[dict, ...]:
        return tuple(
            dict(zip(
                [id_key, vals[0]],
                (chain(item, rep_sex_info, rep_cross_info)
                 for item in items)))
            for items in zip(ids, vals[1:]))
    return (non_transposed, transposed)

def file_data(zfile: ZipFile,
              member_key: str,
              cdata: dict,
              process_value: Optional[Callable[[dict], dict]] = None,
              process_transposed_value: Optional[Callable[
                  [str, tuple[str, ...], tuple[str, ...]],
                  tuple[dict, ...]]] = None) -> Iterator[dict]:
    """Load data from files in R/qtl2 zip bundle."""
    def __default_process_value_non_transposed__(val: dict) -> dict:
        return {
            key: replace_na_strings(cdata, value) for key,value in val.items()
        }

    def __default_process_value_transposed__(
            id_key: str,
            ids: tuple[str, ...],
            vals: tuple[str, ...]) -> tuple[dict, ...]:
        """Default values processor for transposed files."""
        return tuple(
            dict(zip([id_key, replace_na_strings(cdata, vals[0])], items))
            for items in zip(
                    ids, (replace_na_strings(cdata, val) for val in vals[1:])))

    process_value = process_value or __default_process_value_non_transposed__
    process_transposed_value = (
        process_transposed_value or __default_process_value_transposed__)

    try:
        if isinstance(cdata[member_key], list):
            for row in (line for lines in
                        (file_data(
                            zfile, member_key, {**cdata, member_key: innerfile},
                            process_value, process_transposed_value)
                         for innerfile in cdata[member_key])
                        for line in lines):
                yield row
            return
        if not cdata.get(f"{member_key}_transposed", False):
            for row in with_non_transposed(zfile, member_key, cdata, process_value):
                yield row
            return

        for row in with_transposed(
                zfile, member_key, cdata, process_transposed_value):
            yield row
    except KeyError as exc:
        raise MissingFileError(*exc.args) from exc

def cross_information(zfile: ZipFile, cdata: dict) -> Iterator[dict]:
    """Load cross information where present."""
    cdata_cross_info = cdata.get("cross_info", {})
    cross_info_file_key = "covar"
    new_cdata = {**cdata}
    sex_fields = (cdata.get("sex",{}).get("covar",""),)
    if "file" in cdata_cross_info:
        cross_info_file_key = "gnqc_cross_info_file"
        new_cdata = {**cdata, "gnqc_cross_info_file": cdata_cross_info["file"]}

    for row in file_data(zfile,
                         cross_info_file_key,
                         new_cdata,
                         *make_process_data_covar(cdata)):
        yield {
            key: chain(value, partial(replace_cross_info, cdata=cdata))
            for key, value in row.items() if key not in sex_fields}

def sex_information(zfile: ZipFile, cdata: dict) -> Iterator[dict]:
    """Load cross information where present."""
    cdata_sex_info = cdata.get("sex", {})
    sex_info_file_key = "covar"
    new_cdata = {**cdata}
    ci_fields = (cdata.get("cross_info",{}).get("covar",""),)
    if "file" in cdata_sex_info:
        sex_info_file_key = "gnqc_sex_info_file"
        new_cdata = {**cdata, "gnqc_sex_info_file": cdata_sex_info["file"]}

    for row in file_data(zfile,
                         sex_info_file_key,
                         new_cdata,
                         *make_process_data_covar(cdata)):
        yield {
            key: chain(value, partial(replace_sex_info, cdata=cdata))
            for key, value in row.items() if key not in ci_fields}

def genotype_data(zfile: ZipFile):
    """Convenience function to genotype data from R/qtl2 bundle."""
    cdata = control_data(zfile)
    return file_data(zfile, "geno", cdata, *make_process_data_geno(cdata))

def raw_file_data(zipfilepath: Union[str, Path],
                  memberfilename: str) -> Iterator[str]:
    """Read the raw text from a file in the R/qtl2 bundle."""
    with (ZipFile(str(zipfilepath), "r") as zfile,
          zfile.open(memberfilename) as innerfile):
        wrappedfile = io.TextIOWrapper(innerfile)
        for  line in wrappedfile:
            yield line

def strip_comments(rawdata: Iterator[str], commentchar) -> Iterator[str]:
    """Remove comments from raw text."""
    return (line for line in rawdata if not line.startswith(commentchar))

def missing_value_codes_to_none(value: str,
                                nastrings: tuple[str, ...]) -> Optional[str]:
    """
    If 'value' is a missing value code, return `None`, otherwise return 'value'.
    """
    return value if value not in nastrings else None

def replace_genotype_codes(value: str, genocodes: dict):
    """Convert genotype codes into values specified in control file."""
    return genocodes.get(value, value)

def read_control_file(zipfilepath: Union[str, Path]) -> dict:
    """Read control data."""
    with ZipFile(str(zipfilepath), "r") as zfile:
        # move `control_data` code here and replace existing function.
        cdata = control_data(zfile)
        return {
            **cdata,
            **{
                ftype: ([cdata[ftype]]
                        if isinstance(cdata[ftype], str)
                        else cdata[ftype])
                for ftype in FILE_TYPES
                if bool(cdata.get(ftype))
            }
        }


def read_file_data(
        zipfilepath: Union[str, Path],
        memberfilename: str,
        processfile: Callable[[Iterator[str]], Iterator[str]] = lambda itr: itr,
        processline: Callable[[str], str] = lambda line: line,
        processfield: Callable[
            [Optional[str]], Optional[str]] = lambda val: val) -> Iterator[
                tuple[Optional[str], ...]]:
    """Read a single file from the bundle processing each field."""
    cdata = read_control_file(zipfilepath)
    return (
        tuple(processfield(field.strip())
              for field in processline(row.strip()).split(cdata["sep"]))
        for row in
        processfile(
            strip_comments(
                raw_file_data(zipfilepath, memberfilename),
                cdata["comment.char"])))


def read_geno_file_data(
        zipfilepath: Union[str, Path],
        memberfilename: str) -> Iterator[tuple[Optional[str], ...]]:
    """Read a 'geno' file from the R/qtl2 bundle."""
    cdata = read_control_file(zipfilepath)
    return read_file_data(
        zipfilepath,
        memberfilename,
        processfield=partial(
            replace_genotype_codes, genocodes=cdata.get("genotypes", {})))


def load_samples(
        zipfilepath: Union[str, Path], filetype: str) -> tuple[str, ...]:
    """Load the samples/cases/individuals from file(s) of type 'filetype'."""
    cdata = read_control_file(zipfilepath)
    samples: set[str] = set()
    for afile in cdata.get(filetype, []):
        filedata = read_geno_file_data(zipfilepath, afile)
        if cdata.get(f"{filetype}_transposed", False):
            samples.update(
                item for item in next(filedata)[1:] if item is not None)
        else:
            try:
                next(filedata)# Ignore first row.
                samples.update(
                    line[0] for line in filedata if line[0] is not None)
            except StopIteration:# Empty file.
                pass

    return tuple(samples)


load_geno_samples = partial(load_samples, filetype="geno")
load_founder_geno_samples = partial(load_samples, filetype="founder_geno")
load_pheno_samples = partial(load_samples, filetype="pheno")