journal/journal.py

from copy import deepcopy
from datetime import datetime, timedelta
from functools import reduce, partial
from pathlib import Path
from shutil import copyfile, rmtree
from subprocess import run
from tempfile import mktemp, mkdtemp
from zipfile import ZipFile
import json
import random
import re
import sys
import textwrap
import traceback

### GLOBALS

JOURNAL_PATH = Path.home() / '.journal.json'

### UTILS

def compose(*fns):
    return partial(reduce, flip(apply), fns)

def nth_or_default(n, l, default):
    return l[n] if n < len(l) else default

def apply(f, x):
    return f(x)

def flip(f):
    return lambda a1, a2: f(a2, a1)

def identity(x):
    return x

def lazy_get(obj, key, default):
    result = obj.get(key)
    return result if result is not None else default()

def wrap_text(text, columns=80):
    return textwrap.fill(text, columns,
        replace_whitespace=False,
        break_on_hyphens=False,
        break_long_words=False)

def split_keep(delims, string):
    res = []
    buf = []

    def flush_buf():
        nonlocal res, buf
        if buf:
            res.append(''.join(buf))
        buf = []

    for c in string:
        if c in delims:
            flush_buf()
            res.append(c)
        else:
            buf.append(c)

    flush_buf()

    return res

def merge_if(pred, l):
    res = []

    for i, curr in enumerate(l):
        prev = l[i-1] if i-1 >= 0 else None
        if prev and pred(prev, curr):
            res[-1] += curr
        else:
            res.append(curr)

    return res

def editor(fpath):
    run(['nvim', '+', str(fpath)])

def edit(text, suffix=''):
    fpath = Path(mktemp(suffix=suffix))

    fpath.write_text(text)

    editor(fpath)

    text = fpath.read_text()

    fpath.unlink()
    
    return text

def prompt(text):
    return input(text + ' [y/n] ') == 'y'

### DATE UTILS

def parse_date(date):
    return datetime.strptime(date, '%Y-%m-%d')

def format_date(date):
    return date.strftime('%Y-%m-%d')

def parse_timestamp(timestamp):
    return datetime.strptime(timestamp, '%Y-%m-%d %H:%M:%S')

def format_timestamp(timestamp):
    return datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')

def evaluate_time_expression(expression):
    if expression == 'today':
        return datetime.now()
    elif expression == 'yesterday':
        return datetime.now() - timedelta(days=1)

### FILE PARSERS

def parse_foods_file(text):
    foods_str, recipes_str = text.split('---')

    def parse_macro(macro):
        name, value = macro.split()
        return (name, float(value.removesuffix('g').removesuffix('kcal')))

    foods = {
        macros[0]: dict(parse_macro(macro) for macro in macros[1:])
        for macros in [food.split('\n') for food in foods_str.strip().split('\n\n')]
    }

    recipes = {}

    def evaluate_ingredients(ingredients):
        result = {}

        total_weight = 0.0
        for ingredient in ingredients:
            k,v = parse_macro(ingredient)
            if k == 'TOTAL':
                result[k] = v
                continue
            else:
                total_weight += v

            food = foods.get(k)
            total = 100.0
            if not food:
                food = recipes[k].copy()
                total = food['TOTAL']
                del food['TOTAL']

            for kk,vv in food.items():
                if kk not in result:
                    result[kk] = 0.0

                result[kk] += vv * (v/total)

        if 'TOTAL' not in result:
            result['TOTAL'] = total_weight

        return result

    for ingredients in [recipe.split('\n') for recipe in recipes_str.strip().split('\n\n')]:
        recipes[ingredients[0]] = evaluate_ingredients(ingredients[1:])

    return foods, recipes

def evaluate_food_entry(foods, recipes, value, name):
    if name in recipes:
        food = recipes[name]

        if value == 0.0:
            value = food['TOTAL']

        food = {k: v*(value/food['TOTAL']) for k,v in food.items()}
    elif name in foods:
        if value == 0.0:
            value = 100

        food = {k: v*(value/100.0) for k,v in foods[name].items()}
    else:
        breakpoint()
        print(f'ERROR: Invalid diet entry: {name}')
        assert False

    return food

def parse_tasks_file(text):
    result = []

    for task in text.splitlines():
        days, name = task.split(':')
        days = days.split(',')
        result.append((days, name))

    return result

### HACKY HACKERY

_GLOBAL_DO_NOT_USE = {}

def init_hacky_hackery(journal):
    global _GLOBAL_DO_NOT_USE
    _GLOBAL_DO_NOT_USE = journal

def get_foods_file():
    return parse_foods_file(_GLOBAL_DO_NOT_USE['files']['foods'])

### HEADER MODULES

def create_godword(journal, date):
    words = journal['files']['godword'].strip().split('\n')
    return [random.choice(words) for _ in range(20)]

def parse_godword(block):
    return block.split()

def generate_godword(value):
    return f"{' '.join(value[:10])}\n{' '.join(value[10:])}"

def create_habits(journal, date):
    return {x: False for x in journal['files']['habits'].strip().split('\n')}

def parse_habits(block):
    result = {}
    for habit in block.splitlines():
        value, name = habit.split(maxsplit=1)
        result[name.strip()] = value[1] == 'x'
    return result

def generate_habits(value):
    return '\n'.join(f'[{"x" if v else "-"}] {k}' for k,v in value.items())

def create_notifications(journal, date):
    notifications = []

    for day in journal['days'].values():
        for entry in day.get('entries'):
            for block in entry['blocks']:
                if not isinstance(block, str) and block['type'] == 'notify':
                    if block['day'] == date:
                        notifications.append({
                            'source': entry['timestamp'],
                            'message': block['message']
                        })

    return notifications

def parse_notifications(notifications):
    result = []
    for notification in notifications.splitlines():
        parts = notification.split()
        result.append({
            'source': int(parse_timestamp(' '.join(parts[0:2]).strip('[]')).timestamp()),
            'message': ' '.join(parts[2:]),
        })
    return result

def generate_notifications(value):
    return '\n'.join(f'[[{format_timestamp(n["source"])}]] {n["message"]}' for n in value)

def create_tasks(journal, date):
    tasks = parse_tasks_file(journal['files']['tasks'])

    curr_day = {
        0: 'mo', 1: 'tu', 2: 'we', 3: 'th',
        4: 'fr', 5: 'sa', 6: 'su',
    }[parse_date(date).weekday()]

    return {name: False for days, name in tasks if curr_day in days}

def parse_tasks(tasks):
    result = {}
    for task in tasks.splitlines():
        value, name = task.split(maxsplit=1)
        name = name.strip()
        result[name] = value[1] == 'x'
    return result

def generate_tasks(value):
    return '\n'.join(f'[{"x" if v else "-"}] {k}' for k,v in value.items())

def create_sticky(journal, date):
    yesterday = format_date(parse_date(date) - timedelta(days=1))

    if day := journal['days'].get(yesterday):
        if sticky := day['header'].get('sticky'):
            return sticky

### ENTRY MODULES

def parse_timer(block):
    rest = block.split()

    name = None
    timestamp = None
    if len(rest) > 2:
        name, *rest = rest
    if len(rest) > 1:
        timestamp = int(parse_timestamp(' '.join(rest)).timestamp())

    result = {}
    if name:
        result['name'] = name
    if timestamp:
        result['timestamp'] = timestamp

    return result

def generate_timer(block):
    parts = []

    if name := block.get('name'):
        parts.append(name)

    if ts := block.get('timestamp'):
        parts.append(format_timestamp(ts))

    return ' '.join(parts)

def parse_exercise(block):
    parts = block.split()

    if parts[0] == 'walk':
        kind, minutes, distance, steps = parts
        return {
            'kind': kind,
            'minutes': int(minutes.removesuffix('min')),
            'distance': float(distance.removesuffix('km')),
            'steps': int(steps.removesuffix('steps')),
        }
    
    elif parts[0] == 'calisthenics':
        kind, split, exercise = parts
        sets, reps = split.split('x')
        return {
            'kind': kind,
            'reps': reps,
            'sets': sets,
            'exercise': exercise,
        }

    assert False

def generate_exercise(block):
    if block['kind'] == 'walk':
        return f'walk {block["minutes"]}min {block["distance"]}km {block["steps"]}steps'
    elif block['kind'] == 'calisthenics':
        return f'calisthenics {block["sets"]}x{block["reps"]} {block["exercise"]}'

    assert False

DEFAULT_PARSER = lambda b: {'value': b}
DEFAULT_GENERATOR = lambda b: b['value']

entry_modules = {
    'diet': (
        lambda b: {'amount': int(b.split()[0].removesuffix('g')), 'food': b.split()[1].strip()},
        lambda b: f'{b["amount"]}g {b["food"]}'),
    'exercise': (parse_exercise, generate_exercise),
    'behavior': (DEFAULT_PARSER, DEFAULT_GENERATOR),

    'hide': (lambda _: {}, lambda _: ''),
    'info': (DEFAULT_PARSER, compose(DEFAULT_GENERATOR, wrap_text)),

    'post': (
        lambda b: {'timestamp': int(parse_timestamp(b.removeprefix('@post ').strip()).timestamp())},
        lambda b: format_timestamp(b["timestamp"])
    ),
    'notes': (
        compose(
            lambda b: b.splitlines(),
            lambda s: {'source': s[0], 'title': s[1]},
        ),
        lambda b: f'{b["source"]}\n{b["title"]}'
    ),

    'task': (DEFAULT_PARSER, DEFAULT_GENERATOR),
    'start': (parse_timer, generate_timer),
    'stop': (parse_timer, generate_timer),
    'done': (parse_timer, generate_timer),

    'notify': (
        compose(
            lambda b: b.split(maxsplit=1),
            lambda s: {'day': s[0], 'message': s[1]}
        ),
        lambda b: f'{b["day"]} {b["message"]}'
    ),
}

def parse_entry_module(block: str) -> dict:
    tag = block.split()[0].removeprefix('@')
    block = block.removeprefix(f'@{tag}').strip()

    return {'type': tag} | entry_modules[tag][0](block)

def generate_entry_module(block: dict) -> str:
    if block['type'] == 'notes':
        return f'@notes\n{entry_modules[block["type"]][1](block)}'

    return f'@{block["type"]} {entry_modules[block["type"]][1](block)}'

### READ-ONLY STATS SECTION FUNCTIONS

def generate_stats(page):
    if not page['entries']:
        return ''
    
    result = ''

    num_entries = len(page['entries'])
    num_blocks = sum(len(entry['blocks']) for entry in page['entries'])
    text_concat = ' '.join(b for e in page['entries'] for b in e['blocks'] if isinstance(b, str))
    num_words = len(text_concat.split())
    result += f'Entries: {num_entries}, Blocks: {num_blocks}, Words: {num_words}'

    last_entry = max(e['timestamp'] for e in page['entries'])
    first_entry = min(e['timestamp'] for e in page['entries'])
    entry_delta = last_entry - first_entry
    entry_hours = round(entry_delta / 60 / 60, 2)
    result += f'\nFirst: {format_timestamp(first_entry)}, Last: {format_timestamp(last_entry)}, Hours: {entry_hours}'

    calories = 0
    carbs = 0
    fat = 0
    protein = 0
    sugar = 0
    num_meals = 0
    first_meal = float('inf')
    last_meal = float('-inf')

    foods, recipes = get_foods_file()
    for entry in page['entries']:
        did_count = False
        for block in entry['blocks']:
            if not isinstance(block, str) and block['type'] == 'diet':
                food = evaluate_food_entry(foods, recipes, block['amount'], block['food'])

                if not did_count:
                    num_meals += 1
                    first_meal = min(entry['timestamp'], first_meal)
                    last_meal = max(entry['timestamp'], last_meal)
                    did_count = True
                
                calories += food['Energy']
                carbs += food.get('Carbs', 0)
                fat += food.get('Fat', 0)
                protein += food.get('Protein', 0)
                sugar += food.get('Sugar', 0)
        
    carbs_proportion = round(carbs * 4 / calories * 100) if carbs and calories else 0
    fat_proportion = round(fat * 9 / calories * 100) if fat and calories else 0
    protein_proportion = round(protein * 4 / calories * 100) if protein and calories else 0
    
    calories, carbs, fat, protein, sugar = map(partial(round, ndigits=2), [calories, carbs, fat, protein, sugar])

    meal_delta = last_meal - first_meal
    meal_hours = round(meal_delta / 60 / 60, 2)

    result += f'\nCalories: {calories} ({carbs_proportion}/{fat_proportion}/{protein_proportion}, {protein}/{sugar}), Meals: {num_meals}, Hours: {meal_hours}'

    return result

### PAGE FUNCTIONS

def create_header(journal, date):
    return {
        'godword': create_godword(journal, date),
        'habits': create_habits(journal, date),
        'notifications': create_notifications(journal, date),
        'tasks': create_tasks(journal, date),
        'sticky': create_sticky(journal, date),
    }

def parse_header(header):
    header_modules = {
        'godword': parse_godword,
        'habits': parse_habits,
        'notifications': parse_notifications,
        'tasks': parse_tasks,
        'sticky': identity,
    }

    def split_into_blocks(text):
        return [b.strip() for b in re.split(r'\n{2,}', text) if b.strip() != '']

    modules = split_into_blocks(header)

    result = {}

    for module in modules:
        name, value = module.split('\n', maxsplit=1)
        name = name.lower().removesuffix(':')
        result[name] = header_modules[name](value)
    
    return result

def generate_header(header):
    header_modules = {
        'godword': generate_godword,
        'habits': generate_habits,
        'notifications': generate_notifications,
        'tasks': generate_tasks,
        'sticky': identity,
    }

    result = ''

    for name, value in header.items():
        if not value:
            continue

        result += f'\n\n{name.title()}:\n'
        result += header_modules[name](value)
    
    return result

def create_entry(journal, date):
    return {
        'timestamp': int(datetime.now().timestamp()),
        'blocks': []
    }

def parse_entry(entry):
    def merge_notes_block(l):
        res = []

        i = 0
        while i < len(l):
            if l[i] == '@notes':
                # notes nl source nl title
                res.append('\n'.join([l[i], l[i+2], l[i+4]]))
                i += 5
            else:
                res.append(l[i])
                i += 1

        return res

    def merge_wrapped_lines(l):
        TIMESTAMP_LENGTH = len('2020-02-02 02:02:02 ')
        COLUMN_LIMIT = 80

        res = []

        i = 0
        while i < len(l):
            curr = l[i]
            prev = l[i-1] if i > 0 else None
            next = l[i+1] if i+1 < len(l) else None

            # ['aoeu', '\n', 'aoeu']
            if prev and curr == '\n' and next:
                len_prev = len(prev)

                # first block is preceded by timestamp
                if i - 1 == 0:
                    len_prev += TIMESTAMP_LENGTH

                # do not wrap indented lines
                if not next[0].isspace():
                    next_word = next.split()[0]

                    # merge only if text is actually wrapped
                    if len_prev + len(next_word) >= COLUMN_LIMIT:
                        res[-1] += ' ' + next
                        i += 2
                        continue

            res.append(curr)
            i += 1

        return res

    def split_post_block(l):
        res = []

        POST_BLOCK_LENGTH = len('@post 2020-02-02 02:02:02')

        i = 0
        while i < len(l):
            curr = l[i]

            if curr.startswith('@post'):
                res.append(curr[:POST_BLOCK_LENGTH])
                res.append(curr[POST_BLOCK_LENGTH+1:])
            else:
                res.append(curr)

            i += 1
        
        return res

    split_into_blocks = compose(
        # split the text into sections by newline and tag symbol, keeping the separators
        partial(split_keep, ('\n', '@')), 

        # merge sequential newlines together into a single whitespace block 
        partial(merge_if, lambda p, c: p == c == '\n'),

        ## TAG PARSING
        # attach escaped tags
        partial(merge_if, lambda p, c: c == '@' and p[-1] == '\\'), 
        # attach tag
        partial(merge_if, lambda p, c: p == '@'),
        # attach tags which do not come after newline or another tag
        partial(merge_if, lambda p, c: c[0] == '@' and not (not p[-1] != '\n' or (p[0] == '@' and p[-1] == ' '))), 

        ## SPECIAL BLOCK PARSING
        # merge notes block (because it spans 3 lines or 5 blocks)
        merge_notes_block,
        # split post block (because next block could be attached to it)
        split_post_block,

        # strip all non-whitespace blocks
        partial(map, lambda s: s if s.isspace() else s.rstrip()), list,
        # merge escaped tags with following text
        partial(merge_if, lambda p, c: p.endswith('\\@')), 

        # merge wrapped lines
        merge_wrapped_lines,

        # remove trailing whitespace block
        lambda b: b if b and not all(c == '\n' for c in b[-1]) else b[:-1],
    )

    timestamp, content = entry

    return {
        'timestamp': int(parse_timestamp(timestamp.strip()).timestamp()),
        'blocks': [parse_entry_module(b) if b.startswith('@') else b for b in split_into_blocks(content)],
    }

def generate_entry(entry):
    def format_block(block, is_first):
        def format_text(text):
            if all(c == '\n' for c in block):
                return text

            DUMMY_TS = '2020-02-02 02:02:02 '

            if is_first:
                text = DUMMY_TS + text

            length = len(text)
            
            if length > 80:
                text = wrap_text(text)
            
            if is_first:
                text = text.removeprefix(DUMMY_TS)
            
            return text
        
        formatted = format_text(block) if isinstance(block, str) else generate_entry_module(block)

        if result[-1] != '\n' and not all(c == '\n' for c in formatted):
            formatted = ' ' + formatted
        
        return formatted

    result = f'\n\n{format_timestamp(entry["timestamp"])}'
        
    for i, block in enumerate(entry['blocks']):
        result += format_block(block, i == 0)
    
    return result

def create_day(journal, date):
    return {
        'title': date,
        'header': create_header(journal, date),
        'entries': []
    }

def parse_day(text):
    # discard read-only QS section
    text = text[text.find('#'):]

    ENTRY_RE = re.compile(r'^(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}) ?', re.MULTILINE)

    header, *tmp = ENTRY_RE.split(text)
    entries = list(zip(tmp[::2], tmp[1::2]))

    title, *header = header.split('\n', maxsplit=1)
    header = header[0] if len(header) else ''

    return {
        'title': title.removeprefix('# '),
        'header': parse_header(header),
        'entries': [parse_entry(e) for e in entries],
    }

def generate_day(page):
    result = ''

    result += generate_stats(page)

    result += f'\n\n# {page["title"]}'

    result += generate_header(page['header'])

    for entry in page['entries']:
        result += generate_entry(entry)

    result += '\n'

    return result

### COMMAND UTILS

def import_journal(fpath):
    result = { 'days': {}, 'files': {} }

    for name in list(sorted(fpath.glob('*.md'))):
        day = parse_day(name.read_text())
        result['days'][name.stem] = day

    for fname in ['habits', 'godword', 'tasks', 'foods']:
        result['files'][fname] = (fpath / fname).read_text()

    return result

def export_journal(journal, fpath):
    for day in journal['days'].values():
        (fpath / (day['title'] + '.md')).write_text(generate_day(day))
    
    for fname, content in journal['files'].items():
        (fpath / fname).write_text(content)
    
def backup_journal():
    print('Creating backup...')

    tmpdir = mkdtemp()

    journal = load_journal()

    export_journal(journal, Path(tmpdir))
    copyfile(str(JOURNAL_PATH), tmpdir + '/journal.json')

    files = Path(tmpdir).glob('*')

    current_date = datetime.now().strftime('%Y-%m-%d')

    zipfile_path = Path(f'{current_date}.zip')

    zipfile = ZipFile(zipfile_path, 'w')
    
    for file in files:
        zipfile.write(file, arcname=file.name)

    rmtree(tmpdir)

    if script := journal['files'].get('backup'):
        print('Found script, running...')
        run(['bash', '-c', script.replace('%ZIPFILE%', f"'{str(zipfile_path.absolute())}'")])

        if prompt('Ran script, delete archive?'):
            zipfile_path.unlink()

def open_journal(date):
    journal = load_journal()

    if not journal['days'].get(date):
        backup_journal()
        journal['days'][date] = create_day(journal, date)

    tmpdir = Path(mkdtemp())
    export_journal(journal, tmpdir)

    while True:
        try:
            editor(tmpdir / f'{date}.md')
            new_journal = import_journal(tmpdir)
            break
        except Exception as e:
            traceback.print_exc()
            input('Press enter to try again...')

    save_journal(new_journal)

def load_journal():
    if JOURNAL_PATH.exists():
        return json.loads(JOURNAL_PATH.read_text())
    else:
        return import_journal(Path.home() / 'workspace' / 'journal')

def save_journal(journal):
    JOURNAL_PATH.write_text(json.dumps(journal))

### COMMAND HANDLERS

def handle_open(args):
    subcommand = nth_or_default(0, args, 'today')

    if date := evaluate_time_expression(subcommand):
        open_journal(format_date(date))
    else:
        print(f'Invalid subcommand: {subcommand}')

def handle_edit(args):
    subcommand = nth_or_default(0, args, 'foods')

    journal = load_journal()

    if subcommand in journal['files']:
        journal['files'][subcommand] = edit(journal['files'][subcommand])
    elif prompt(f'Unknown file: {subcommand}, create new?'):
        journal['files'][subcommand] = edit('')

    save_journal(journal)

def handle_import(args):
    if len(args) < 1:
        print('Missing directory.')
        return

    path = Path(args[0])

    if not path.is_dir():
        print(f'Invalid directory: {path}')
        return

    save_journal(import_journal(path))

def handle_export(args):
    if len(args) < 1:
        print('Missing directory.')
        return
    
    path = Path(args[0])

    if not path.is_dir():
        print(f'Invalid directory: {path}')
        return
    
    export_journal(load_journal(), path)

def handle_test(args):
    journal = load_journal()

    journal_orig = deepcopy(journal)

    for day in journal['days']:
        journal['days'][day] = parse_day(generate_day(journal['days'][day]))

    if journal != journal_orig:
        print('Test failed!')

        print('Dumping journal.fail.json and journal.fail.orig.json...')
        Path('journal.fail.json').write_text(json.dumps(journal, indent=4))
        Path('journal.fail.orig.json').write_text(json.dumps(journal_orig, indent=4))
    else:
        print('Test passed!')


def handle_summary(args):
    def generate_food_summary(page):
        result = ''

        def print(str=''):
            nonlocal result
            result += '\n' + str
        
        foods, recipes = get_foods_file()

        daily_calories = 0.0
        daily_protein = 0.0

        for entry in page['entries']:
            has_printed = False
            entry_calories = 0.0
            entry_protein = 0.0
            for diet in (b for b in entry['blocks'] if type(b) != str and b['type'] == 'diet'):
                if not has_printed:
                    print(f'-- {format_timestamp(entry["timestamp"])}')
                    has_printed = True

                value = diet['amount']
                name = diet['food']

                if name in recipes:
                    food = recipes[name]

                    if value == 0.0:
                        value = food['TOTAL']

                    food = {k: v*(value/food['TOTAL']) for k,v in food.items()}
                elif name in foods:
                    if value == 0.0:
                        value = 100

                    food = {k: v*(value/100.0) for k,v in foods[name].items()}
                else:
                    print(f'ERROR: Invalid diet entry: {diet}')
                    continue

                protein = round(food.get('Protein', 0.0), 2)
                calories = round(food.get('Energy', 0.0), 2)

                entry_calories += calories
                entry_protein += protein

                print(f'{name:<20} {value:<6}g, {calories:<6}kcal, {protein:<6}g protein')

            if has_printed:
                entry_calories = round(entry_calories, 2)
                entry_protein = round(entry_protein, 2)
                print(f'-- TOTAL: {entry_calories}kcal, {entry_protein}g protein')
                print()

            daily_calories += entry_calories
            daily_protein += entry_protein

        print(f'-- DAILY TOTAL ({daily_calories}kcal, {daily_protein}g protein)')

        return result

    subcommand = nth_or_default(0, args, 'today')

    date = evaluate_time_expression(subcommand)
    if not date:
        print(f'Invalid time expression: {subcommand}')
        return
    
    date = format_date(date)

    journal = load_journal()

    print(generate_food_summary(journal['days'][date]))

### MAIN

def main():
    init_hacky_hackery(load_journal())

    command = nth_or_default(1, sys.argv, 'open')
    args = sys.argv[2:]

    handler = {
        'open': handle_open,
        'edit': handle_edit,
        'import': handle_import,
        'export': handle_export,
        'test': handle_test,
        'summary': handle_summary,
        'backup': lambda _: backup_journal(),
    }.get(command, lambda _: print(f'Invalid command: {command}'))

    handler(args)

if __name__ == '__main__':
    main()