#! /usr/bin/env python3

import argparse
import numpy as np
import re
import sys

# Character sets for sequence types. Noise symbol X/N needs to be final character.
charset = {
    "DNA": [c for c in "ACGTN"],
    "PEPTIDE": [c for c in "ACDEFGHIKLMNPQRSTVWYX"],
    "RNA": [c for c in "ACGUN"],
}

parser = argparse.ArgumentParser(epilog="Specify SEED to generate the same sequences.")
parser.add_argument(
    "-L",
    "--label",
    default=False,
    help="prefix for the sequence ids.",
)
parser.add_argument(
    "-l",
    "--len",
    type=int,
    default=100,
    help="mean length of output sequences, default 100.",
)
parser.add_argument(
    "-n",
    "--num",
    type=int,
    default=1,
    help="number of sequences to generate, default 1.",
)
parser.add_argument(
    "-N", "--noise", type=float, default=0.0, help="noise percentage, default 0.0."
)
parser.add_argument(
    "-r",
    "--rand",
    action="store_true",
    default=False,
    help="sequences have random length, Poisson-distributed with mean LEN.",
)
parser.add_argument(
    "-s",
    "--seed",
    type=int,
    default=False,
    help="random number generator seed.",
)
parser.add_argument(
    "-t",
    "--type",
    choices=["DNA", "RNA", "PEPTIDE"],
    default="DNA",
    help="type of sequence(s) to generate.",
)
args = parser.parse_args()

prefix = ""
if args.label:
    tmp = re.sub(r"\W+", "_", args.label.strip())
    if tmp:  # skip if args.label was all whitespace.
        prefix = tmp + ":"

if args.noise < 0 or args.noise > 100:
    print(
        f"[ERROR] noise percentage needs to be between 0 and 100, specified {args.noise}",
        file=sys.stderr,
    )
    exit(1)
if args.len < 1:
    print(
        f"[ERROR] target sequence length needs to be greater than 0, specified {args.len}",
        file=sys.stderr,
    )
    exit(1)

pn = args.noise / 100.0
po = (1 - pn) / (len(charset[args.type]) - 1)
probs = np.full(len(charset[args.type]), po)
probs[-1] = pn  # final position is the noise symbol

if args.seed:
    rng = np.random.default_rng(np.random.PCG64(args.seed))
else:
    rng = np.random.default_rng()

for i in range(args.num):
    if args.rand:
        ##seqlen = max(1, np.random.poisson(args.len))
        seqlen = max(1, rng.poisson(args.len))
    else:
        seqlen = args.len

    ##chars = np.random.choice(charset[args.type], size=seqlen, p=probs,)
    chars = rng.choice(
        charset[args.type],
        size=seqlen,
        p=probs,
    )
    print(f">{prefix}{i + 1}")
    print(f"{''.join(chars)}")
