import hashlib
import json
import math
import os
import random
import re
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from .lora import LoRA, file_sha256
from .pipeline import plain, prefix_mask

AUDIO = ('.wav', '.flac', '.mp3', '.ogg', '.opus', '.aiff', '.aif')
OTHER_AUDIO = ('.m4a', '.aac', '.wma', '.mp4', '.webm', '.mov')
SPEECH_TOKENS = 'speech_tokenizer_v3_25hz'
REF_WHOLE, REF_EXCERPT = 400, 375
SAMPLE_TEXT = dict(en='Hello there! This is a quick test of my new voice. How does it sound to you?',
                   ja='こんにちは。これは新しい声のテストです。どう聞こえますか？',
                   mix='Hello! これは新しい声のテストです。How does it sound?')
CHARS_PER_SECOND = dict(en=((2.5, 35), (6, 24)), ja=((1, 25), (3, 15)), mix=((1, 35), (3, 24)))

PROBLEMS = {
    'no_text': 'no transcript (add a .txt file with the same name)',
    'empty_text': 'the transcript is empty',
    'unreadable': "the audio can't be read",
    'other_format': "this audio format can't be read (convert it to .wav or .flac)",
    'low_rate': 'recorded below 16 kHz (too low to learn a voice from)',
    'silent': 'almost silent',
    'short': 'too short',
    'long': 'too long (split it into shorter clips)',
    'mismatch': "the transcript doesn't fit the length of the audio (check the .txt)",
    'tokens': "the audio couldn't be processed",
}
WARNINGS = {
    'rate': 'recorded below 44.1 kHz: the voice will learn a duller top end',
    'clipped': 'clipped (distorted peaks)',
    'digits': 'digits in the transcript: write numbers out as words, the way they are spoken',
    'pace': 'the transcript may not fit the audio (unusually fast or slow speech for its length)',
    'duplicate': 'the same transcript appears more than once',
    'orphan': 'a transcript with no audio file of the same name',
}


def frames_db(x, rate):
    n = max(1, rate // 100); k = len(x) // n
    if k == 0: return np.full(1, -120.)
    return 20 * np.log10(np.sqrt((x[:k * n].reshape(k, n) ** 2).mean(1)) + 1e-9)


def trim(x, rate):
    db = frames_db(x, rate); loud = float(db.max())
    on = np.where(db > max(min(-45., loud - 30), -60.))[0]
    if len(on) == 0: return x[:0], loud
    n = rate // 100
    return x[max(0, on[0] * n - int(.15 * rate)):min(len(x), (on[-1] + 1) * n + int(.25 * rate))], loud


def read_mono(path):
    x, rate = sf.read(str(path), dtype='float32', always_2d=True)
    return x.mean(1), rate


def scan(folder, lang='en', min_seconds=1., max_seconds=30.):
    folder = Path(folder)
    if not folder.is_dir(): raise SystemExit(f'{folder} is not a folder')
    desc_file = folder/'voice.txt'
    description = plain(' '.join(desc_file.read_text(encoding='utf-8').split())) if desc_file.exists() else None
    problems, warnings, clips, seen = defaultdict(list), defaultdict(list), [], defaultdict(list)
    files = sorted(p for p in folder.rglob('*') if p.is_file() and not any(part.startswith('.') for part in p.relative_to(folder).parts))
    hard, soft = CHARS_PER_SECOND.get(lang, CHARS_PER_SECOND['en'])
    for p in files:
        rel = p.relative_to(folder).as_posix(); ext = p.suffix.lower()
        if ext in OTHER_AUDIO: problems['other_format'].append(rel); continue
        if ext not in AUDIO: continue
        txt = p.with_suffix('.txt')
        if not txt.exists(): problems['no_text'].append(rel); continue
        text = plain(' '.join(txt.read_text(encoding='utf-8', errors='replace').split()))
        if not text: problems['empty_text'].append(rel); continue
        try:
            x, rate = read_mono(p)
        except Exception:
            problems['unreadable'].append(rel); continue
        if rate < 16000: problems['low_rate'].append(rel); continue
        y, loud = trim(x, rate); seconds = len(y) / rate
        if loud < -40 or seconds == 0: problems['silent'].append(rel); continue
        if seconds < min_seconds: problems['short'].append(f'{rel} ({seconds:.1f} s)'); continue
        if seconds > max_seconds: problems['long'].append(f'{rel} ({seconds:.0f} s)'); continue
        cps = len(re.sub(r'\[[^\]]*\]|[\W_]', '', text)) / seconds
        if not hard[0] <= cps <= hard[1]: problems['mismatch'].append(rel); continue
        if not soft[0] <= cps <= soft[1]: warnings['pace'].append(rel)
        if rate < 44100: warnings['rate'].append(rel)
        if (np.abs(x) >= .999).mean() > 1e-3: warnings['clipped'].append(rel)
        if re.search(r'\d', re.sub(r'\[[^\]]*\]', '', text)): warnings['digits'].append(rel)
        seen[text.lower()].append(rel)
        st = p.stat()
        key = hashlib.sha1(f'{rel}|{st.st_size}|{st.st_mtime_ns}|{text}|{min_seconds}|{max_seconds}'.encode()).hexdigest()[:20]
        clips.append(dict(file=rel, text=text, seconds=round(seconds, 2), rate=rate, key=key))
    for rels in seen.values():
        if len(rels) > 1: warnings['duplicate'] += rels
    audio = {p.with_suffix('') for p in files if p.suffix.lower() in AUDIO + OTHER_AUDIO}
    warnings['orphan'] += [p.relative_to(folder).as_posix() for p in files
                           if p.suffix.lower() == '.txt' and p != desc_file and p.with_suffix('') not in audio]
    if not warnings['orphan']: del warnings['orphan']
    return dict(clips=clips, problems=problems, warnings=warnings, description=description, folder=folder)


def describe(found, log=print):
    clips, problems, warnings = found['clips'], found['problems'], found['warnings']
    left = sum(len(v) for v in problems.values())
    minutes = sum(c['seconds'] for c in clips) / 60
    log(f"Checked {len(clips) + left} clips in {found['folder']}")
    log(f'  ready: {len(clips)} clips, {minutes:.1f} minutes of speech')
    if left:
        log(f'  left out: {left} clips')
        for kind, rels in problems.items():
            log(f'    {PROBLEMS[kind]}: {len(rels)}')
            for r in rels[:8]: log(f'      {r}')
            if len(rels) > 8: log(f'      ... and {len(rels) - 8} more')
    for kind, rels in warnings.items():
        log(f'  note, {WARNINGS[kind]}: {len(rels)}' + (f' ({", ".join(rels[:3])}{"..." if len(rels) > 3 else ""})' if len(rels) < len(clips) else ''))
    log(f'  voice description (voice.txt): {found["description"]!r}' if found['description'] else
        '  voice description: none (optional: describe the voice in a voice.txt file in the folder)')
    if minutes < 2:
        raise SystemExit(f'Only {minutes:.1f} minutes of usable speech. A LoRA needs at least 2 minutes; 20 to 60 minutes works best.')
    if minutes < 10:
        log('  note: under 10 minutes of speech. It can work, but 20 to 60 minutes gives a much closer voice.')


@torch.no_grad()
def encode(robin, found, cache, log=print, stop=None, progress=None):
    from torchaudio.functional import resample
    cache = Path(cache); cache.mkdir(parents=True, exist_ok=True)
    vae_id = file_sha256(robin.vae_file) if getattr(robin, 'vae_file', None) else 'module'
    meta_file = cache/'meta.json'; meta = dict(vae=vae_id, speech_tokens=SPEECH_TOKENS)
    if meta_file.exists() and json.loads(meta_file.read_text()) != meta:
        for f in cache.glob('*.npz'): f.unlink()
    meta_file.write_text(json.dumps(meta))
    out, bad, t0, done = [], [], time.time(), 0
    todo = sum(not (cache/f"{c['key']}.npz").exists() for c in found['clips'])
    if todo: log(f'Preparing {todo} clips...')
    for c in found['clips']:
        f = cache/f"{c['key']}.npz"
        if not f.exists():
            if stop is not None and stop():
                raise SystemExit(f'Stopped while preparing clips ({done} of {todo} done; they are kept for next time).')
            try:
                x, rate = read_mono(found['folder']/c['file'])
                y = torch.from_numpy(trim(x, rate)[0])
                tokens = robin.speech(y.numpy(), rate)
                z = robin.vae.encode(resample(y, rate, 48000)[None].repeat(2, 1)[None].to(robin.device))[0].T.float().cpu().numpy()
                if abs(len(tokens) - len(z)) > 2: raise ValueError('length mismatch')
                k = min(len(tokens), len(z))
                np.savez(f, tokens=tokens[:k].astype(np.int16), latents=z[:k].astype(np.float16))
            except Exception:
                bad.append(c['file']); continue
            done += 1
            if progress is not None: progress('prepare', done, todo)
            if done % 25 == 0 or done == todo: log(f'  {done}/{todo} prepared ({time.time() - t0:.0f} s)')
        d = np.load(f)
        out.append(dict(c, tokens=d['tokens'].astype(np.int64), latents=d['latents']))
    if bad:
        found['problems']['tokens'] += bad
        log(f"  left out {len(bad)} clips: {PROBLEMS['tokens']}: {', '.join(bad[:5])}{'...' if len(bad) > 5 else ''}")
    return out


class Trainer:
    def __init__(self, robin, clips, description, out, *, lang='en', steps=1000, rank_planner=32, rank_decoder=64, lr_planner=5e-5,
                 lr_decoder=2e-4, accum=8, warmup=100, lr_floor=.2, p_reference=.4, save_every=200, sample_text=None, seed=0, log=print, stop=None,
                 progress=None):
        self.robin, self.description, self.out, self.lang, self.log, self.stop = robin, description, Path(out), lang, log, stop
        self.progress = progress
        self.hp = dict(steps=steps, rank_planner=rank_planner, rank_decoder=rank_decoder, lr_planner=lr_planner, lr_decoder=lr_decoder,
                       accum=accum, warmup=min(warmup, max(1, steps // 10)), lr_floor=lr_floor, p_reference=p_reference, save_every=save_every,
                       lang=lang, seed=seed)
        self.sample_text = sample_text or SAMPLE_TEXT.get(lang, SAMPLE_TEXT['en'])
        rng = random.Random(seed); order = list(range(len(clips))); rng.shuffle(order); torch.manual_seed(seed)
        hold = min(10, len(clips) // 20) if len(clips) >= 40 else 0
        self.holdout = [clips[i] for i in order[:hold]]; self.clips = [clips[i] for i in order[hold:]]
        self.short = [c for c in self.clips if len(c['tokens']) <= REF_WHOLE]
        ranks = {k: r for k, r in (('planner', rank_planner), ('decoder', rank_decoder)) if r > 0}
        nets = dict(planner=robin.planner.layers, decoder=robin.decoder.model.layers)
        self.lora = LoRA({k: nets[k] for k in ranks}, ranks, seed=seed).attach()
        groups = [dict(params=[p for n, p in zip(self.lora.names, getattr(self.lora, which)) if n.startswith(net + '.')],
                       lr=dict(planner=lr_planner, decoder=lr_decoder)[net]) for net in ranks for which in ('A', 'B')]
        for g in groups: g['base_lr'] = g['lr']
        self.optimizer = torch.optim.AdamW(groups, betas=(.9, .95), weight_decay=0., eps=1e-8)
        self.rng, self.step, self.norms = rng, 0, []
        self.sel = torch.cat((torch.arange(robin.offset, robin.offset + 32768), torch.tensor([robin.audio_end]))).to(robin.device)
        self.base_sha = robin.base_sha256()

    def lr_scale(self, step):
        hp = self.hp
        if step <= hp['warmup']: return step / hp['warmup']
        f = min(1., (step - hp['warmup']) / max(1, hp['steps'] - hp['warmup']))
        return hp['lr_floor'] + (1 - hp['lr_floor']) * .5 * (1 + math.cos(math.pi * f))

    def example(self, clip, rng):
        mode = 'reference' if len(self.clips) > 1 and rng.random() < self.hp['p_reference'] else 'tts'
        desc = self.description if mode == 'tts' or rng.random() < .5 else None
        if mode == 'tts' and rng.random() < .1: desc = None
        ref = ref_text = None
        if mode == 'reference':
            pool = [c for c in (self.short or self.clips) if c is not clip] or [c for c in self.clips if c is not clip]
            other = rng.choice(pool); tok = other['tokens']
            if len(tok) <= REF_WHOLE:
                ref, ref_text = tok, other['text']
            else:
                s = rng.randrange(len(tok) - REF_EXCERPT + 1); ref = tok[s:s + REF_EXCERPT]
        return self.robin.prompt(clip['text'], desc, self.lang, ref, ref_text)

    def losses(self, ids, clip, train=True, gen=None, sigmas=None):
        r = self.robin; dev = r.device
        seq = torch.tensor(list(ids) + [int(t) + r.offset for t in clip['tokens']] + [r.audio_end], device=dev)
        ss, n = len(ids), len(seq)
        mask = prefix_mask(n, ss, dev)
        h = r.planner(seq[None], mask)[0][ss - 1:-1]
        targets = seq[ss:]
        targets = torch.where(targets == r.audio_end, torch.full_like(targets, 32768), targets - r.offset)
        W = r.planner.lm_head.weight[self.sel]
        ce = h.new_zeros((), dtype=torch.float32)
        for s in range(0, len(targets), 256):
            ce = ce + checkpoint(lambda hs, ys: F.cross_entropy(F.linear(hs, W).float(), ys, reduction='sum'),
                                 h[s:s + 256], targets[s:s + 256], use_reentrant=False)
        ce = ce / len(targets)
        if train: (ce / self.hp['accum']).backward()
        del h
        with torch.no_grad():
            ks, vs, _ = r.planner.prefill(seq, mask)
        clean = torch.from_numpy(clip['latents'].astype(np.float32)).to(dev)
        flows, sigmas = [], sigmas or [min(.98, max(.02, self.rng.betavariate(2, 2))) for _ in range(2)]
        for sigma in sigmas:
            noise = torch.randn(clean.shape, device=dev, generator=gen)
            xt = sigma * noise + (1 - sigma) * clean
            v = r.decoder(xt[None], torch.tensor([sigma], device=dev), ks, vs, checkpointing=train)[0]
            fl = F.mse_loss(v, noise - clean)
            if train: (fl / (len(sigmas) * self.hp['accum'])).backward()
            flows.append(float(fl.detach()))
        return float(ce.detach()), sum(flows) / len(flows)

    @torch.no_grad()
    def evaluate(self):
        if not self.holdout: return None
        gen = torch.Generator(device=self.robin.device).manual_seed(1)
        rows = [self.losses(self.robin.prompt(c['text'], self.description, self.lang), c, train=False, gen=gen, sigmas=(.25, .5, .75))
                for c in self.holdout]
        return dict(planner=round(sum(r[0] for r in rows) / len(rows), 4), decoder=round(sum(r[1] for r in rows) / len(rows), 4))

    @torch.no_grad()
    def sample(self, name):
        if self.robin.vae is None: return
        try:
            wav, sr, _ = self.robin.speak(self.sample_text, self.description, lang=self.lang, seed=0)
            (self.out/'samples').mkdir(exist_ok=True)
            path = self.out/'samples'/f'{name}.wav'
            sf.write(str(path), wav, sr, subtype='PCM_16')
        except Exception as e:
            self.log(f'  (could not render a sample: {e})')
            return
        if self.progress is not None: self.progress('sample', self.step, self.hp['steps'], path=path)

    def save(self, final=False, light=False):
        (self.out/'checkpoints').mkdir(exist_ok=True)
        path = self.out/'checkpoints'/f'step_{self.step:05d}.safetensors'
        self.lora.save(path, self.base_sha)
        if final: self.lora.save(self.out/'lora.safetensors', self.base_sha)
        state = dict(step=self.step, hp=self.hp, lora=self.lora.state_dict(), optimizer=self.optimizer.state_dict(), rng=self.rng.getstate(),
                     torch_rng=torch.get_rng_state(), cuda_rng=torch.cuda.get_rng_state() if torch.cuda.is_available() else None, norms=self.norms)
        tmp = self.out/'resume.pt.tmp'; torch.save(state, tmp); os.replace(tmp, self.out/'resume.pt')
        ev = None if light else self.evaluate()
        if not light: self.sample(f'step_{self.step:05d}')
        self.log(f'  saved {path.relative_to(self.out)}' + (f' | held-out loss: planner {ev["planner"]:.3f}, decoder {ev["decoder"]:.4f}' if ev else ''))
        with (self.out/'log.jsonl').open('a') as f: f.write(json.dumps(dict(step=self.step, saved=path.name, heldout=ev)) + '\n')

    def resume(self):
        state = torch.load(self.out/'resume.pt', map_location='cpu', weights_only=False)
        keep = ('rank_planner', 'rank_decoder', 'lang')
        if any(state['hp'][k] != self.hp[k] for k in keep):
            raise SystemExit(f'{self.out} holds a run with different settings ({", ".join(f"{k}={state["hp"][k]}" for k in keep)}). '
                             'Use the same settings, a new output folder, or start over (--restart).')
        self.lora.load_state_dict(state['lora']); self.optimizer.load_state_dict(state['optimizer'])
        self.rng.setstate(state['rng']); torch.set_rng_state(state['torch_rng'])
        if state.get('cuda_rng') is not None and torch.cuda.is_available(): torch.cuda.set_rng_state(state['cuda_rng'])
        self.step, self.norms = state['step'], state.get('norms', [])

    def run(self):
        hp, r = self.hp, self.robin
        self.out.mkdir(parents=True, exist_ok=True)
        (self.out/'train_config.json').write_text(json.dumps(dict(hp, clips=len(self.clips), held_out=len(self.holdout),
                                                                  description=bool(self.description)), indent=1))
        if self.step == 0:
            self.lora.scale = 0; self.sample('base'); self.lora.scale = 1
        params = [p for p in self.lora.parameters()]
        n_params = sum(p.numel() for p in params)
        self.log(f'Training a LoRA ({n_params / 1e6:.0f}M parameters) on {len(self.clips)} clips'
                 + (f', {len(self.holdout)} held out for checking' if self.holdout else '') + f', {hp["steps"]} steps.')
        if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats()
        t_start, step_times, run_loss = time.time(), [], []
        while self.step < hp['steps']:
            t0 = time.time(); self.step += 1
            ces, flows = [], []
            for _ in range(hp['accum']):
                for _attempt in range(4):
                    clip = self.rng.choice(self.clips)
                    try:
                        ce, fl = self.losses(self.example(clip, self.rng), clip)
                        ces.append(ce); flows.append(fl); break
                    except torch.cuda.OutOfMemoryError:
                        self.optimizer.zero_grad(set_to_none=True); torch.cuda.empty_cache(); ces, flows = [], []
                        self.log(f"  ran out of GPU memory on {clip['file']} ({clip['seconds']:.0f} s), leaving it out. "
                                 'If this keeps happening, use a smaller --max-seconds.')
                    except Exception as e:
                        self.optimizer.zero_grad(set_to_none=True); ces, flows = [], []
                        self.log(f"  couldn't train on {clip['file']} ({type(e).__name__}: {e}), leaving it out.")
                    self.clips = [c for c in self.clips if c is not clip]
                    self.short = [c for c in self.short if c is not clip]
                    if not self.clips: raise SystemExit('No clips left to train on.')
            for g in self.optimizer.param_groups: g['lr'] = g['base_lr'] * self.lr_scale(self.step)
            norm = float(torch.nn.utils.clip_grad_norm_(params, 1.0))
            med = sorted(self.norms)[len(self.norms) // 2] if len(self.norms) >= 50 else None
            if not ces or not math.isfinite(norm) or (med is not None and norm > 20 * med):
                self.optimizer.zero_grad(set_to_none=True)
                if ces: self.log(f'  step {self.step}: an unusual update (gradient {norm:.2f}) was skipped')
            else:
                self.norms = (self.norms + [norm])[-200:]
                self.optimizer.step(); self.optimizer.zero_grad(set_to_none=True)
            step_times.append(time.time() - t0); run_loss.append((np.mean(ces) if ces else float('nan'), np.mean(flows) if flows else float('nan')))
            with (self.out/'log.jsonl').open('a') as f:
                f.write(json.dumps(dict(step=self.step, planner=round(float(run_loss[-1][0]), 4), decoder=round(float(run_loss[-1][1]), 5),
                                        grad=round(norm, 4), seconds=round(step_times[-1], 2))) + '\n')
            pace = float(np.mean(step_times[-20:])); left = (hp['steps'] - self.step) * pace
            pl, dl = np.nanmean([x[0] for x in run_loss[-10:]]), np.nanmean([x[1] for x in run_loss[-10:]])
            if self.progress is not None:
                self.progress('train', self.step, hp['steps'], planner=float(pl), decoder=float(dl), pace=pace, left=left)
            if self.step % 10 == 0 or self.step <= 3 or self.step == hp['steps']:
                mem = f' | GPU {torch.cuda.max_memory_allocated() / 2 ** 30:.1f} GB' if torch.cuda.is_available() else ''
                self.log(f"step {self.step}/{hp['steps']} | planner loss {pl:.3f} | decoder loss {dl:.4f} | {pace:.1f} s/step | "
                         f"{left / 60:.0f} min left{mem}")
            if self.step % hp['save_every'] == 0 or self.step == hp['steps']:
                self.save(final=self.step == hp['steps'])
            elif self.stop is not None and self.stop():
                self.save(light=True)
            if self.stop is not None and self.stop() and self.step < hp['steps']:
                self.log(f'Stopped at step {self.step}. Train again with the same folder to continue from here.')
                return None
        self.log(f"Done in {(time.time() - t_start) / 60:.0f} min. The LoRA is {self.out/'lora.safetensors'}; every saved step is in "
                 f"{self.out/'checkpoints'} with a sample of each in {self.out/'samples'}.")
        return self.out/'lora.safetensors'


def train(robin, data, out, *, lang='en', min_seconds=1., max_seconds=30., restart=False, log=print, stop=None, progress=None, **settings):
    out = Path(out)
    if restart:
        old = [out/f for f in ('resume.pt', 'lora.safetensors', 'log.jsonl')] + list(out.glob('checkpoints/step_*.safetensors')) + list(out.glob('samples/*.wav'))
        for f in old:
            if f.exists(): f.unlink()
    elif (out/'lora.safetensors').exists() and not (out/'resume.pt').exists():
        raise SystemExit(f'{out} already holds a finished LoRA. Use a new output folder, or start over (--restart) to train again.')
    found = scan(data, lang, min_seconds, max_seconds); describe(found, log)
    clips = encode(robin, found, out/'cache', log, stop, progress)
    robin._speech = None
    if torch.cuda.is_available(): torch.cuda.empty_cache()
    trainer = Trainer(robin, clips, found['description'], out, lang=lang, log=log, stop=stop, progress=progress, **settings)
    try:
        if (out/'resume.pt').exists():
            trainer.resume()
            if trainer.step >= trainer.hp['steps']:
                if not (out/'lora.safetensors').exists(): trainer.lora.save(out/'lora.safetensors', trainer.base_sha)
                log(f'This run already finished ({trainer.step} steps). Start over (--restart) to train again, or give more steps (--steps) to keep going.')
                return out/'lora.safetensors'
            log(f'Continuing the unfinished run from step {trainer.step}.')
        return trainer.run()
    finally:
        trainer.lora.detach()