#!/usr/bin/env python3.6

# pylint: disable=line-too-long,no-self-use

"""Script for finding cancer hot spots. Running this script with flag will run the
whole process from scratch.  In case some of the components fail in the middle,
running this script again will pick the process up from where it failed.

Examples
--------

./frag_hot_spots.py --cancer=lung_cancer --policy=diff:2000 --policy=diff:10000

./frag_hot_spots.py --cancer=lung_cancer --skip-aggr --skip-sort --policy=diff:2000 --policy=diff:10000

python3 bclassify_1d_data.py --data ./frag-2018-04-09-lung_cancer-end-754a55dc0caa9948-3ff692f8999e910d-diff:5000-features.tsv --labels ~/tsv/solid+liquid_binary_labels.csv --n_comp=80 --mod_z_threshold=4 --drop_pattern='_[ML]'


Adding a new target selection policy
------------------------------------


The HotSpotSelectionPolicy class implements a policy to pick "interesting" hot
spots given the hot spotms from Y&H and cancer cohorts. If just want to
experiment with a new policy, do the following:

- Create a new policy object with a new, unique name.
- Add the policy to the --policy flag and plumb its handling.
- Run

  ./frag_hot_spots.py --no_run_aggr --no_run_sort --no_run_features --policy=newpolicy

  where 'newpolicy' is the name of the policy you just implemented.
  This run should finish quickly (about one minute), and it produces a BED file in:

  $GRAIL/go/src/grail.com/cmd/grail-query/.results/frag-somethingorother-bed.

  Check this BED file and verify that it makes sense.  This BED file has 7
  fields.

  5th field is the ranking of this location, using
  (long frag freq + short frag freq) as the sort criterion.

  6th field is the ranking of this location, using short frag freq as the sort
  criterion.

  7th field is the ranking of this location, using long frag freq as the sort
  criterion.

- Run

  ./frag_hot_spots.py --no_run_aggr --no_run_sort --policy=newpolicy


  This run should take about an hour.

  It creates file ./frag-somethingorother-features.tsv.  You can feed this file
  as the --data file to bclassify_1d_data.


For more details

https://docs.google.com/document/d/1GPsNmrg6GgSDKtaPpWWuj1IVhvduOK5_g4jLLM1qn0M/edit

https://docs.google.com/document/d/13NKE0cRrMo8ClcQ8o8UaHUIz3IC7ggpJGm7jLMbZ6ws/edit#heading=h.aui10pydrd7h

"""

import argparse
import concurrent.futures
import enum
import logging
import os
import sys
import itertools
from typing import Callable, Dict, Iterable, List, Set, Tuple, Optional, NamedTuple

import scipy.stats.mstats
import gql
import ccga

@enum.unique
class FragLength(enum.Enum):
    """Fragment length type"""
    SHORT = 1  # read fragments whose length < 155
    LONG = 2   # read fragments whose length >= 165
    def __str__(self) -> str:
        if self.value == 1:
            return 'short'
        return 'long'

@enum.unique
class Endpoint(enum.Enum):
    """Fragment endpoint type."""
    START = 1  # Aggregate fragments by their start positions
    END = 2    # Aggregate fragments by their end positions
    def __str__(self) -> str:
        if self.value == 1:
            return 'start'
        return 'end'

    @staticmethod
    def parse(name: str) -> 'Endpoint':
        "Inverse of Endpoint.__str__."
        for endpoint in [Endpoint.END, Endpoint.START]:
            if name == endpoint.__str__():
                return endpoint
        raise Exception(f'Endpoint {name} not found')

Path = str # Alias for readability.

ROOT = 's3://grail-query/results'

def location_frequency_path(config: 'Config',
                            cohort: ccga.Cohort,
                            length: FragLength,
                            endpoint: Endpoint,
                            fold: int) -> str:
    """File that stores genomic location -> frequency mapping"""
    patient_subset = config.kfold.training(cohort, fold)
    return f'{ROOT}/frag-{config.date}-{fold}-of-{ccga.MAX_FOLDS}- {cohort}-{length}-{endpoint}-{patient_subset.hash_str}-locs.btsv'

def top100k_locations_path(config: 'Config',
                           cohort: ccga.Cohort,
                           length: FragLength,
                           endpoint: Endpoint,
                           fold: int) -> str:
    """File that stores top 100k hot locations. It is generated by sorting
    location_frequency_path file"""
    patient_subset = config.kfold.training(cohort, fold)
    return f'{ROOT}/frag-{config.date}-{fold}-of-{ccga.MAX_FOLDS}-{cohort}-{length}-{endpoint}-{patient_subset.hash_str}-locs.btsv-top100k.tsv'

class HotLocation:
    """Information about one genomic location (chr, pos) for one cohort (e.g., cancer patients)."""
    def __init__(self, ref: str, pos: int, name: str) -> None:
        """
        Args:
        ref: chromosome name, e.g., "chr1".
        pos: zero-based position within the ref.
        name: string identifying (ref, pos). In practice it's "<ref>_<pos>.
        It is generated in pick_top100k().
        """
        # Normalized frequency of long DNA fragments starting/ending at this location.
        self.long = 0.0
        # Normalized frequency of short DNA fragments starting/ending at this location.
        self.short = 0.0
        # Reference name, e.g., "chr1"
        self.ref = ref
        # Position in ref. Zero based.
        self.pos = pos
         # Mnemonic, e.g., "chr1_12345"
        self.name = name

        # Scaled self.long value so that sum of the long field values across the
        # genome adds up to 1.0
        self.long_normalized = 0.0
        # Scaled self.short value so that sum of the short field values across
        # the genome adds up to 1.0
        self.short_normalized = 0.0

        # Following fields indicate the ranking of this location (0-based)
        # within the entire genome.  Filled by HotLocations.finalize.
        self.long_rank = -1
        self.short_rank = -1
        self.sum_rank = -1

ZERO_LOCATION = HotLocation('zero', -1, 'zero')

class HotLocations:
    """HotLocations accumulates information about hot locations.

    It is logically a map from locname -> HotLocation, where locname is
    HotLocations.name
    """

    def __init__(self) -> None:
        # Keys are location names, such as "chr1_12345"
        self.__locs: Dict[str, HotLocation] = {}

    def add_long(self, ref: str, pos: int, name: str, freq: float):
        """Register the fact that long fragments are observed at the given location."""
        loc = self.__intern(ref, pos, name)
        loc.long += freq

    def add_short(self, ref: str, pos: int, name: str, freq: float):
        """Register the fact that short fragments are observed at the given location."""
        loc = self.__intern(ref, pos, name)
        loc.short += freq

    def finalize(self) -> None:
        """Finalize should be called after adding locations via add_{long,short}.  This
        method computes and fills derived stats. No more add* calls shall be
        made in a future.

        """

        sum_short, sum_long = 0.0, 0.0
        for loc in sorted(self.__locs.values(), key=lambda loc: loc.long):
            sum_long += loc.long
        for loc in sorted(self.__locs.values(), key=lambda loc: loc.short):
            sum_short += loc.short
        for loc in self.__locs.values():
            loc.long_normalized = loc.long / sum_long
            loc.short_normalized = loc.short / sum_short

        rank = 0
        for loc in sorted(self.__locs.values(), key=lambda loc: -loc.long):
            loc.long_rank = rank
            rank += 1

        rank = 0
        for loc in sorted(self.__locs.values(), key=lambda loc: -loc.short):
            loc.short_rank = rank
            rank += 1

        rank = 0
        for loc in sorted(self.__locs.values(), key=lambda loc: -(loc.short+loc.long)):
            loc.sum_rank = rank
            rank += 1

    def find(self, name: str) -> HotLocation:
        """Find the frequency of long/short fragments at the given location. Returns
        ZERO_LOCATION if not found."""
        return self.__locs.get(name, ZERO_LOCATION)

    def find_or_none(self, name: str) -> Optional[HotLocation]:
        """Find the frequency of long/short fragments at the given location. Returns
        None if not found."""
        return self.__locs.get(name, None)

    def find_or_die(self, name: str) -> HotLocation:
        """Find the frequency of long/short fragments at the given location. Dies
        if not found."""
        return self.__locs[name]

    def locs(self) -> Iterable[HotLocation]:
        """List all registered locations."""
        for value in self.__locs.values():
            yield value

    def loc_names(self) -> Iterable[str]:
        """List all registered locations."""
        for value in self.__locs.values():
            yield value.name

    def __intern(self, ref: str, pos: int, name: str) -> HotLocation:
        if name not in self.__locs:
            self.__locs[name] = HotLocation(ref, pos, name)
        return self.__locs[name]

def merged_hot_locations(gql_sess: gql.GQL, long_path, short_path: Path) -> HotLocations:
    """Read the topN TSV files for long and short fragments and merge them."""

    locs = HotLocations()
    for path in (long_path, short_path):
        s: Set[str] = set()
        for row in gql_sess.open_tsv(path):
            ref = row[0]
            pos = int(row[1])
            name = row[2]
            freq = float(row[3])
            if name in s:
                logging.error('Duplicate %s in %s', name, path)
            s.add(name)
            if path == long_path:
                locs.add_long(ref, pos, name, freq)
            else:
                locs.add_short(ref, pos, name, freq)
    locs.finalize()
    return locs

GQLResult = NamedTuple(
    'GQLResult',
    [('tsv', Optional[Path]),
     ('error', Optional[str])])

def aggregate_frag_endpoints(config: 'Config',
                             cohort: ccga.Cohort,
                             length: FragLength,
                             endpoint: Endpoint,
                             fold: int) -> GQLResult:
    """Find all the *.prio files that match the condition, and create a btsv file
    mapping genomic locations to fragment frequencies. It returns the path of
    the generated btsv file.

    date: tidy release date, e.g., "2018-04-09".
    cohort: Sample population.
    length: fragments to read.
    endpoint: START (END) aggregagates fragments by their start (end, resp) points.

    """

    patient_subset = config.kfold.training(cohort, fold)
    out_path = location_frequency_path(config, cohort, length, endpoint, fold)
    logging.info('Aggregating to %s', out_path)

    endpoint_expr = '$start'
    if endpoint == Endpoint.END:
        endpoint_expr = '$start+$length-1'
    length_filter = '$length >= 165'
    if length == FragLength.SHORT:
        length_filter = '$length < 155'

    expr = f"""
out_path := `{out_path}`;
frag_files := {cohort}_frag_files({config.date});
subpatients := read(`{patient_subset.tsv_path}`);
subset_frag_files := join({{frag: frag_files, subpatients: subpatients}}, frag.patient_id==subpatients.patient_id, map:={{frag./.*/}});

locs := map(subset_frag_files, map(yh.prio,
   {{$reference, $start, $length, sample_id: yh.sample_id, scale_factor: yh.scale_factor}},
   filter:=!$duplicate && {length_filter}), row:=yh);

print(sprintf("Reading %d fragment tables", count(locs)));

flatten(locs, subshard:=true)
  | reduce({{$reference, pos:{endpoint_expr}}},
          float(_acc) + float(_val),
          map:=1.0 / $scale_factor,
          shards:=2000)
  | write(out_path, shards:=64);
"""
    label = f'aggr_{cohort}_{length}_{endpoint}'
    status = config.gql.run(expr=expr, label=label)
    if status != 0:
        return GQLResult(tsv=None, error=f'aggregate {label}: process failed with status {status}')
    return GQLResult(tsv=out_path, error=None)

def pick_top100k(config: 'Config',
                 in_tsvs: List[str]) -> List[str]:
    """For each file listed in the input, find top 100k hottest genomic locations
    and dump them in a new TSV file.  The input tsv should be generated by
    aggregate_frag_endpoints().  Returns the pathnames of the generated TSV
    files, one for each in_tsv.

    """

    expr = """
pick_top100k := func(in, out) {
  read(in)
  | filter(string_has_prefix($key.reference, "chr"))
  | minn(100000, -$value, shards:=2000)
  | map({reference: $key.reference, pos: $key.pos, name: sprintf("%s_%d", $key.reference, $key.pos), count: $value})
  | write(out);
};

"""

    out_paths: List[str] = []
    for in_path in in_tsvs:
        out_path = f'{ROOT}/{os.path.basename(in_path)}-top100k.tsv'
        logging.info('top100k: %s->%s', in_path, out_path)
        expr += f'pick_top100k(`{in_path}`, `{out_path}`);\n'
        out_paths.append(out_path)

    label = 'top100k'
    status = config.gql.run(expr=expr, label=label)
    if status != 0:
        raise Exception(f'top100k {label}: process failed with status {status}')
    return out_paths

def build_feature_table(config: 'Config',
                        label: str,
                        bed_path: str,
                        cancer_cohort: ccga.Cohort,
                        endpoint: Endpoint,
                        fold: int) -> GQLResult:
    """Scan all the *.prio files in CFDNA dataset, find fragments that intersect the
    tragets in bed_path."""

    cancer_test_patients = config.kfold.test(cancer_cohort, fold)
    healthy_test_patients = config.kfold.test(ccga.Cohort.HEALTHY, fold)

    file_name_prefix = f'frag-{config.date}-{fold}-of-{ccga.MAX_FOLDS}-{cancer_cohort}-{endpoint}-{cancer_test_patients.hash_str}-{healthy_test_patients.hash_str}-{label}'
    out_path = f'{ROOT}/{file_name_prefix}-features.btsv'

    endpoint_expr = '$start'
    if endpoint == Endpoint.END:
        endpoint_expr = '$start+$length-1'

    # Read a fragment (.prio) file and intersect with a BED file.
    expr = f"""
quantize_length := func(n) {{ cond(n >= 165,  "L", cond(n < 155, "S", "M")) }};

bed := read(`{bed_path}`);

read_prio := func(sample_id, prio, bed, num_fragments_in_file) {{
  checkpoint_path := sprintf("s3://grail-query/checkpoint/{file_name_prefix}-%s.tsv", sample_id);

  prio
  | joinbed(bed, chrom:=$reference, start:={endpoint_expr}, length:=1, map:={{feat:feat.featname, qlength:quantize_length($length)}})
  | reduce({{$feat, $qlength}}, int(_acc)+int(_val), map:=1)
  | map({{feat:$key.feat, qlength:$key.qlength, count: float($value) / float(num_fragments_in_file)}})
  | write(checkpoint_path);
  print(sprintf("read_prio: wrote %s", checkpoint_path));
  read(checkpoint_path);
}};

frag_files := {cancer_cohort}_and_healthy_frag_files({config.date});
subpatients := flatten(table(read(`{cancer_test_patients.tsv_path}`), read(`{healthy_test_patients.tsv_path}`)));
subset_frag_files := join({{frag: frag_files, subpatients: subpatients}}, frag.patient_id==subpatients.patient_id, map:={{frag./.*/}});

// For every fragments.prio file in the samples, intersect it $bed and write the result in /tmp/filtered_<sampleid>.btsv.
subset_frag_files | map({{sample_id: $sample_id, value: read_prio($sample_id, $prio, bed, $fragment_count)}}, shards:=1024)
                  | transpose({{sample_id: $sample_id}}, {{$feat, $qlength, $count}})
                  | write(`{out_path}`);
"""
    label = f'feature_{cancer_cohort}_{endpoint}_{label}'
    status = config.gql.run(expr=expr, label=label)
    if status != 0:
        GQLResult(tsv=None, error=f'aggregate {label}: process failed with status {status}')

    out_tsv_path = out_path.replace('.btsv', '.tsv')
    config.gql.copy_tsv(out_path, out_tsv_path)
    local_tsv_path = f'./{file_name_prefix}-features.tsv'
    config.gql.copy_file(out_tsv_path, local_tsv_path)
    return GQLResult(tsv=out_path, error=None)

# Scorer is a function that computes the "hotness score" given the gene locus
# info for the cancer cohort and young&healthy cohort.
Scorer = Callable[[HotLocation, HotLocation], float]

class HotSpotSelectionPolicy:
    """HotSpotSelectionPolicy implements a policy of picking the locations to extract from samples.


    The select() method is given a hot locations from Y&H band cancer
    populations. It should return a set of locations to sample.

    """
    def name(self) -> str:
        """Return the name of this policy."""
        raise Exception('not implemented')

    def select(self, cancer_locs: HotLocations, yh_locs: HotLocations) -> Iterable[HotLocation]:
        """Select is called with the cancer and young&healthy pileups. It should return
        the list of interesting genomic locations.

        """
        raise Exception('not implemented')

class UnnormalizedWassersteinPolicy(HotSpotSelectionPolicy):
    """The policy that simply ranks locations a diff between hot cancer locations
    and hot y&h locations"""
    def __init__(self, name: str, n: int) -> None:
        self.__name = name
        self.__n = n

    def name(self) -> str:
        """Implements the base class's name method."""
        return self.__name

    def select(self, cancer_locs: HotLocations, yh_locs: HotLocations) -> Iterable[HotLocation]:
        """Implements the base class's select method."""

        all_loc_names = set(cancer_locs.loc_names()).union(yh_locs.loc_names())
        sort_key = lambda name: -self.score(cancer_locs.find(name), yh_locs.find(name))
        for name in sorted(all_loc_names, key=sort_key)[:self.__n]:
            yield cancer_locs.find_or_none(name) or yh_locs.find_or_die(name)

    def score(self, target: HotLocation, other: HotLocation) -> float:
        """Compute an earth-mover distance between the two fragment distributions.

        https://en.wikipedia.org/wiki/Earth_mover%27s_distance.
        """
        sum_score = 0.0
        emd1 = (target.short - other.short)
        sum_score += abs(emd1)
        emd2 = emd1 + (target.long - other.long)
        sum_score += abs(emd2)
        return sum_score

class ZNormalizedWassersteinPolicy(HotSpotSelectionPolicy):
    """The policy that simply ranks locations a diff between hot cancer locations
    and hot y&h locations"""
    def __init__(self, name: str, n: int) -> None:
        self.__name = name
        self.__n = n

    def name(self) -> str:
        """Implements the base class's name method."""
        return self.__name

    def __total_freq(self, locs0: HotLocations, locs1: HotLocations, name: str) -> float:
        loc0 = locs0.find(name)
        loc1 = locs1.find(name)
        return loc0.long_normalized + loc0.short_normalized + loc1.long_normalized + loc1.short_normalized

    def __z_normalized_freqs(self,
                             cancer_locs: HotLocations,
                             yh_locs: HotLocations) -> Dict[str, float]:
        names = [x for x in set(cancer_locs.loc_names()).union(yh_locs.loc_names())]
        freqs = [self.__total_freq(cancer_locs, yh_locs, name) for name in names]
        z_normalized_freqs = scipy.stats.mstats.zscore(freqs)

        z_normalized_freqs_map: Dict[str, float] = {}
        assert len(names) == len(z_normalized_freqs)
        for i, name in enumerate(names):
            z_normalized_freqs_map[name] = z_normalized_freqs[i]
        return z_normalized_freqs_map

    def select(self, cancer_locs: HotLocations, yh_locs: HotLocations) -> Iterable[HotLocation]:
        """Implements the base class's select method."""

        all_loc_names = set(cancer_locs.loc_names()).union(yh_locs.loc_names())
        z_normalized_freqs = self.__z_normalized_freqs(cancer_locs, yh_locs)
        sort_key = lambda name: -self.score(cancer_locs.find(name),
                                            yh_locs.find(name),
                                            z_normalized_freqs[name])
        for name in sorted(all_loc_names, key=sort_key)[:self.__n]:
            yield cancer_locs.find_or_none(name) or yh_locs.find_or_die(name)

    def score(self, target: HotLocation, other: HotLocation, z_score: float) -> float:
        """Compute an earth-mover distance between the two fragment distributions.

        https://en.wikipedia.org/wiki/Earth_mover%27s_distance.
        """
        sum_score = 0.0
        emd1 = (target.short - other.short)
        sum_score += abs(emd1)
        emd2 = emd1 + (target.long - other.long)
        sum_score += abs(emd2)

        freq = target.long_normalized + target.short_normalized + other.long_normalized + other.short_normalized
        sum_score /= freq
        sum_score *= z_score
        return sum_score

class DiffPolicy(HotSpotSelectionPolicy):
    """The policy that simply takes a diff between hot cancer locations and hot y&h locations"""
    def __init__(self, name: str, n: int) -> None:
        self.__name = name
        self.__n = n

    def name(self) -> str:
        """Implements the base class's name method."""
        return self.__name

    def __pick_interesting_locs(self, target: HotLocations, other: HotLocations, n: int) -> Iterable[str]:
        """Run the equivalemnt of ScoreBasedPolicy to pick n interesting locations that happen often in target but not often in other.

        If target is a cancer cohort, then the other should be y&h, and vice versa.
        """

        diff_locs: Dict[str, Tuple[str, int, float]] = {}
        for target_loc in target.locs():
            other_loc = other.find(target_loc.name)
            diff_locs[target_loc.name] = (
                target_loc.ref,
                target_loc.pos,
                DiffPolicy.__score(target_loc, other_loc))
        i = 0
        for name, unused in sorted(diff_locs.items(), key=lambda v: -v[1][2]):
            yield name
            i += 1
            if i >= n:
                break

    @staticmethod
    def __score(target: HotLocation, other: HotLocation) -> float:
        """Compute the score for a genomic location. The higher score means
        more interesting.
        """

        target_mean = (target.short + target.long) / 2
        other_mean = (other.short + other.long) / 2

        diff2 = (
            max(target_mean - other_mean, 0) +
            abs((target.short - target.long) - (other.short - other.long))
        )
        return diff2

    def select(self, cancer_locs: HotLocations, yh_locs: HotLocations) -> Iterable[HotLocation]:
        """Implements the base class's select method."""
        locs0 = set(self.__pick_interesting_locs(yh_locs, cancer_locs, self.__n))
        locs1 = set(self.__pick_interesting_locs(cancer_locs, yh_locs, self.__n))

        logging.info(f'|locs0|=%d |locs1|=%d', len(locs0), len(locs1))
        diff01 = locs0 - locs1
        diff10 = locs1 - locs0
        logging.info(f'|diff01|=%d |diff10|=%d', len(diff01), len(diff10))

        for name in diff01.union(diff10):
            loc = cancer_locs.find(name)
            if loc == ZERO_LOCATION:
                loc = yh_locs.find(name)
                assert loc != ZERO_LOCATION
            yield loc

def generate_bed(config: 'Config',
                 cancer_cohort: ccga.Cohort,
                 endpoint: Endpoint,
                 fold: int,
                 policy: HotSpotSelectionPolicy) -> str:
    """Given the top100k TSV files for the given date and the endpoint, this
    function computes interesting genomic coordinates to look at. It generates a
    BED file and returns its path.

    """

    cancer_training_patients = config.kfold.training(cancer_cohort, fold)

    cancer_locs = merged_hot_locations(
        config.gql,
        top100k_locations_path(config, cancer_cohort,
                               FragLength.LONG, endpoint, fold),
        top100k_locations_path(config, cancer_cohort,
                               FragLength.SHORT, endpoint, fold))
    yh_locs = merged_hot_locations(
        config.gql,
        top100k_locations_path(config, ccga.Cohort.YOUNG_AND_HEALTHY,
                               FragLength.LONG, endpoint, fold),
        top100k_locations_path(config, ccga.Cohort.YOUNG_AND_HEALTHY,
                               FragLength.SHORT, endpoint, fold))

    local_targets_path = f'{config.result_dir}/frag-{config.date}-{fold}-of-{ccga.MAX_FOLDS}-{cancer_cohort}-{endpoint}-{cancer_training_patients.hash_str}-{policy.name()}-targets.tsv'
    logging.info('creating %s', local_targets_path)
    with open(local_targets_path, 'w') as targets:
        n = 0
        print('chrom\tstart\tend\tfeatname\tsumrank\tshortrank\tlongrank\tshortfreq\tlongfreq', file=targets)
        for loc in policy.select(cancer_locs, yh_locs):
            print(f'{loc.ref}\t{loc.pos}\t{loc.pos+1}\t{loc.name}\t{loc.sum_rank}\t{loc.short_rank}\t{loc.long_rank}\t{loc.short_normalized*1e6}\t{loc.long_normalized*1e6}', file=targets)
            n += 1
            if n > 10000:
                break
    targets_path = f'{ROOT}/{os.path.basename(local_targets_path)}'
    logging.info('creating %s', targets_path)
    config.gql.copy_file(local_targets_path, targets_path)
    return targets_path

class Config:
    """Config parses commandline flags into a more digestible form"""

    def __init__(self) -> None:
        parser = argparse.ArgumentParser()
        parser.add_argument('--date',
                            default='2018-04-09',
                            help='Tidy release date')
        parser.add_argument('--skip-aggr',
                            action='store_true',
                            help='If false, skip the location aggregation step')
        parser.add_argument('--skip-sort',
                            action='store_true',
                            help='If false, skip the location ranking/sorting step')
        parser.add_argument('--skip-features',
                            action='store_true',
                            help='If false, skip the featurevector generation step')
        parser.add_argument('--policy', type=str, default=[], action='append',
                            help="""This flag sets the policy used to produce targets given the hotspot data
                            computed from cohorts (young / cancer), fragment lengths
                            (long/short).

                            Policy name is generally of form type:param. Please read
                            the Config policy parser code to make sense of
                            it. Examples of some good policies include diff:5000.
                            This flag can be set multiple times, and a separate
                            feature TSV file will be generated for each.""")
        parser.add_argument('--cancer', type=str, default=[],
                            action='append',
                            help="""List of cancer types to study.

                            This flag can be specified multiple times.  For example,
                            --cancer=lung --cancer=breast will aggregate length
                            information for lung and breast patients indenpedently,
                            then create a feature file for each. The cancer type is
                            matched against the 'primccat' column in the tidy
                            clinical.tsv table. A special column type 'all' matches
                            all cancer types.""")
        parser.add_argument('--endpoint', type=str, default=[],
                            action='append', choices=['start', 'end'],
                            help="""List of fragment endpoints. By default,
                            data for both endpoints are generated. That is, the default is
                            "--endpoint=start --endpoint=end".
                            """)
        parser.add_argument('--bed', type=str, default=[],
                            action='append',
                            help=""" A BED file to use to create feature TSVs.

                            This flag can be specified up to two times in the command
                            line.  One of the paths must contain word 'start' and the
                            other must contain word 'end'. They are used to create
                            feature TSVs from fragment start and end endpoints,
                            respectively. If this flag is set, --policy is ignored,
                            and --skip-aggr and -skip-sort flags will be all set
                            to true.""")
        parser.add_argument('--overwrite-files',
                            action='store_true',
                            help='Pass --overwrite-files flag to GQL')
        parser.add_argument('--fold', default=[], action='append')
        args = parser.parse_args()

        self.result_dir = os.path.join(os.environ['GRAIL'],
                                       'go/src/grail.com/cmd/grail-query/.results')

        gql_flags: List[str] = []
        self.overwrite_files = args.overwrite_files
        if args.overwrite_files:
            gql_flags.append('--overwrite-files')
        self.gql = gql.GQL(log_prefix='/tmp/frag_hot_spots', default_flags=gql_flags)

        self.run_aggr = not args.skip_aggr
        self.run_sort = not args.skip_sort
        self.run_features = not args.skip_features
        self.date = args.date
        self.policy: List[HotSpotSelectionPolicy] = []
        self.endpoints = [Endpoint.END]
        if args.endpoint:
            self.endpoints = [Endpoint.parse(name) for name in args.endpoint]

        self.folds = [0]
        if args.fold:
            self.folds = []
            for fi in args.fold:
                fold = int(fi)
                if fold < 0 or fold >= ccga.MAX_FOLDS:
                    raise Exception(f'Illegal fold {fold} (must be in range [0,{ccga.MAX_FOLDS})')
                self.folds.append(fold)

        self.policy = [ZNormalizedWassersteinPolicy('z_normalized_wassterstein:6000', 6000)]
        if args.policy:
            self.policy = []
            for name in args.policy:
                v = name.split(':')
                if v[0] == 'diff':
                    n = int(v[1])
                    self.policy.append(DiffPolicy(name, n))
                    continue
                if v[0] == 'unnormalized_wasserstein':
                    n = int(v[1])
                    self.policy.append(UnnormalizedWassersteinPolicy(name, n))
                    continue
                if v[0] == 'z_normalized_wasserstein':
                    n = int(v[1])
                    self.policy.append(ZNormalizedWassersteinPolicy(name, n))
                    continue
                raise Exception(f'Illegal policy {name}')

        self.cancer = [ccga.Cohort.ALL_CANCER]
        if args.cancer:
            self.cancer = [ccga.Cohort.parse(s) for s in args.cancer]

        self.bed: Dict[Endpoint, str] = {}
        if args.bed:
            if not args.skip_aggr:
                logging.info('setting --skip-addr')
                self.run_aggr = False
            if not args.skip_sort:
                logging.info('setting --skip-sort')
                self.run_sort = False

            for path in args.bed:
                if path.find('start') >= 0:
                    if Endpoint.START in self.bed:
                        raise Exception('Multiple start BED files specified')
                    self.bed[Endpoint.START] = path
                if path.find('end') >= 0:
                    if Endpoint.END in self.bed:
                        raise Exception('Multiple start BED files specified')
                    self.bed[Endpoint.END] = path
            if not self.bed:
                raise Exception('--bed must contain at least one start or end bed file')
        self.kfold = ccga.KFold(self.gql, self.date)

def main() -> None:
    """Main entry point.

    Simply running this script with no argument will run the whole pipeline from
    start to finish. If you run the script again, it will pick up from where it
    left off.

    This script must be run in directory $GRAIL/go/src/cmds/grail-query/python.
    """

    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)-15s [%(filename)s:%(lineno)d] %(message)s')
    config = Config()
    logging.getLogger('').addHandler(
        logging.FileHandler(os.path.join(config.gql.log_dir(), 'log.txt')))
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=128)
    logging.info('Commandline: %s', sys.argv)
    logging.info('Storing results in %s, logs in %s', config.result_dir, config.gql.log_dir())
    if config.overwrite_files:
        gql.reset_cache_dir()

    for cohort in config.cancer + [ccga.Cohort.YOUNG_AND_HEALTHY, ccga.Cohort.HEALTHY]:
        for fold in config.folds:
            train = config.kfold.training(cohort, fold)
            test = config.kfold.test(cohort, fold)
            logging.info('%s: training dataset in %s, hash %s, test dataset in %s, hash %s',
                         cohort, train.tsv_path, train.hash_str, test.tsv_path, test.hash_str)
    # Step 1: for each (young&healthy, cancer) * (longfrags, shortfrags) *
    # (start endpoint, end endpoint), generate a btsv file that maps (chr, pos)
    # -> (normalize # of fragments at that location).
    loc_tsvs: List[str] = []
    all_cohorts = config.cancer + [ccga.Cohort.YOUNG_AND_HEALTHY]
    if config.run_aggr:
        gql.send_mail(f'{config.gql.log_dir()}: starting location aggregation', '(eom)')
        results: List[concurrent.futures.Future] = []

        for fold, endpoint, length, cohort in itertools.product(config.folds,
                                                                config.endpoints,
                                                                (FragLength.LONG, FragLength.SHORT),
                                                                all_cohorts):
            logging.info("aggregate: fold=%s, end=%s, length=%s, cohort=%s",
                         fold, endpoint, length, cohort)
            results.append(pool.submit(aggregate_frag_endpoints,
                                       config=config,
                                       cohort=cohort,
                                       length=length,
                                       endpoint=endpoint,
                                       fold=fold))
        all_ok = True
        for r in results:
            result: GQLResult = r.result()
            if result.error:
                logging.error(result.error)
                gql.send_mail(f'{config.gql.log_dir()}: location aggregation failed',
                              'f{result.error}')
                loc_tsvs.append(result.error)
                all_ok = False
                continue
            assert result.tsv
            loc_tsvs.append(result.tsv)
        if not all_ok:
            gql.send_mail(f'{config.gql.log_dir()}: location aggregation failed',
                          f'locations: {loc_tsvs}')
            sys.exit(1)
        gql.send_mail(f'{config.gql.log_dir()}: finished location aggregation',
                      f'locations: {loc_tsvs}')
    else:
        for fold, endpoint, length, cancer_cohort in itertools.product(config.folds,
                                                                       config.endpoints,
                                                                       (FragLength.LONG, FragLength.SHORT),
                                                                       all_cohorts):
            loc_tsvs.append(location_frequency_path(config, cohort, length, endpoint, fold))

    # Step 2: Pick the top100k hottest locations from each file generated in the previous step.
    if config.run_sort:
        gql.send_mail(f'{config.gql.log_dir()}: start top100k ranking', f'(eom)')
        pick_top100k(config, loc_tsvs)
        gql.send_mail(f'{config.gql.log_dir()}: finished top100k ranking', f'(eom)')

    results = []
    for endpoint, cancer_cohort in itertools.product(config.endpoints, config.cancer):
        # Step 3: Generate a BED file from the top100k locations generated in step 2.
        bed_configs: List[Tuple[str, str, int]] = [] # List of (BED path, policy name, fold) to run.
        if config.bed:
            if endpoint in config.bed:
                bed_path = config.bed[endpoint]
                bed_configs.append((bed_path, os.path.basename(bed_path), 9999))
        else:
            for fold in config.folds:
                for policy in config.policy:
                    bed_path = generate_bed(config,
                                            cancer_cohort=cancer_cohort,
                                            endpoint=endpoint,
                                            fold=fold,
                                            policy=policy)
                    bed_configs.append((bed_path, policy.name(), fold))
            if not bed_configs:
                logging.info(f'Skipping %s', endpoint)
                continue

            # Step 4: Intersect the BED and all the CFDNA samples.
            if config.run_features:
                logging.info(f'Generating feature tsvs using configs %s', bed_configs)
                for b in bed_configs:
                    results.append(pool.submit(build_feature_table,
                                               config=config,
                                               label=b[1],
                                               bed_path=b[0],
                                               cancer_cohort=cancer_cohort,
                                               endpoint=endpoint,
                                               fold=b[2]))
    all_ok = True
    for r in results:
        result = r.result()
        if result.error:
            logging.error(result.error)
            gql.send_mail(f'{config.gql.log_dir()}: feature generation failed',
                          f'{result.error}')
            all_ok = False
        assert result.tsv
        logging.info('Generated %s', result.tsv)
    if all_ok:
        gql.send_mail(f'{config.gql.log_dir()}: all done',
                      '(eom)')
    else:
        gql.send_mail(f'{config.gql.log_dir()}: feature generation failed',
                      '(eom)')
        sys.exit(1)

main()
