import os
import torch
import numpy as np
import yaml
import copy
from tqdm import tqdm
from torchaudio.compliance import kaldi
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from fairseq import checkpoint_utils
from transformers import AutoModel, Wav2Vec2FeatureExtractor

from utils.io_optim import (
    TorchaudioDataset,
    LibrosaDataset,
    FFmpegDataset,
    collate_batch,
)
from modules import whisper_extractor as whisper
from modules.wenet_extractor.utils.init_model import init_model
from modules.wenet_extractor.utils.checkpoint import load_checkpoint
from utils.tokenizer import G2PModule, tokenize_text
from utils.symbol_table import SymbolTable
from text.g2p import preprocess_english, read_lexicon

"""
    Extractor for content features
    1. whisper
    2. contentvec
    3. wenet
    4. mert

    Pipeline:
        in preprocess.py:
            call extract_utt_content_features() to extract content features for each utterance
            extract_utt_content_features() envelopes the following steps:
                1. load the model (whisper, contentvec, wenet)
                2. extract the content features
                3. save the content features into files
        in svc_dataset.py:
            call offline_align() to align the content features to the given target length

"""

"""
    Extractor Usage:
        1. initialize an instance of extractor
            extractor = WhisperExtractor(cfg)
        2. load the specified model
            extractor.load_model()
        3. extract the content features
            extractor.extract_content(utt) for single utterance
            extractor.extract_content_batch(utts) for batch utterances
        4. save the content features
            extractor.save_feature(utt, content_feature) for single utterance
"""


class BaseExtractor:
    def __init__(self, cfg):
        self.cfg = cfg
        self.extractor_type = None
        self.model = None

    def offline_align(self, content, target_len):
        """
        args:
            content: (source_len, dim)
            target_len: target length
        return:
            mapped_feature: (target_len, dim)
        """
        target_hop = self.cfg.preprocess.hop_size

        assert self.extractor_type in ["whisper", "contentvec", "wenet"]
        if self.extractor_type == "whisper":
            source_hop = (
                self.cfg.preprocess.whisper_frameshift
                * self.cfg.preprocess.whisper_downsample_rate
                * self.cfg.preprocess.sample_rate
            )
        elif self.extractor_type == "contentvec":
            source_hop = (
                self.cfg.preprocess.contentvec_frameshift
                * self.cfg.preprocess.sample_rate
            )
        elif self.extractor_type == "wenet":
            source_hop = (
                self.cfg.preprocess.wenet_frameshift
                * self.cfg.preprocess.wenet_downsample_rate
                * self.cfg.preprocess.sample_rate
            )
        source_hop = int(source_hop)
        factor = np.gcd(source_hop, target_hop)
        source_hop //= factor
        target_hop //= factor

        # (source_len, 256)
        _, width = content.shape
        # slice the content from padded feature
        source_len = min(target_len * target_hop // source_hop + 1, len(content))

        # const ~= target_len * target_hop
        const = source_len * source_hop // target_hop * target_hop

        # (source_len * source_hop, dim)
        up_sampling_feats = np.repeat(content, source_hop, axis=0)
        # (const, dim) -> (const/target_hop, target_hop, dim) -> (const/target_hop, dim)
        down_sampling_feats = np.average(
            up_sampling_feats[:const].reshape(-1, target_hop, width), axis=1
        )

        err = abs(target_len - len(down_sampling_feats))
        if err > 8:
            # err_log_dir is indeterminate
            err_log_dir = os.path.join(
                self.cfg.preprocess.processed_dir, "align_max_err.log"
            )
            try:
                with open(err_log_dir, "r") as f:
                    err_num = int(f.read())
            except:
                with open(err_log_dir, "w") as f:
                    f.write("0")
                err_num = 0
            if err > err_num:
                with open(err_log_dir, "w") as f:
                    f.write(str(err))

        if len(down_sampling_feats) < target_len:
            # (1, dim) -> (err, dim)
            end = down_sampling_feats[-1][None, :].repeat(err, axis=0)
            down_sampling_feats = np.concatenate([down_sampling_feats, end], axis=0)

        # (target_len, dim)
        mapped_feature = down_sampling_feats[:target_len]

        return mapped_feature

    def save_feature(self, utt, content_feature):
        """Save a single utternace to path {cfg.preprocess.processed_dir}

        Args:
            utt (dict): one item in metadata, containing information for one utterance
            content_feature (tensor): content feature of one utterance
        """
        uid = utt["Uid"]
        assert self.extractor_type != None
        out_dir = os.path.join(
            self.cfg.preprocess.processed_dir, utt["Dataset"], self.extractor_type
        )
        os.makedirs(out_dir, exist_ok=True)
        save_path = os.path.join(out_dir, uid + ".npy")
        # only keep effective parts
        duration = utt["Duration"]
        if self.extractor_type == "whisper":
            frameshift = (
                self.cfg.preprocess.whisper_frameshift
                * self.cfg.preprocess.whisper_downsample_rate
            )  # 20ms
        elif self.extractor_type == "contentvec":
            frameshift = self.cfg.preprocess.contentvec_frameshift  # 20ms
        elif self.extractor_type == "wenet":
            frameshift = (
                self.cfg.preprocess.wenet_frameshift
                * self.cfg.preprocess.wenet_downsample_rate
            )  # 40ms
        elif self.extractor_type == "mert":
            frameshift = self.cfg.preprocess.mert_frameshift
        else:
            raise NotImplementedError
        # calculate the number of valid frames
        num_frames = int(np.ceil((duration - frameshift) / frameshift)) + 1
        # (num_frames, dim) -> (valid_frames, dim)
        assert (
            len(content_feature.shape) == 2
        ), "content feature shape error, it should be (num_frames, dim)"
        content_feature = content_feature[:num_frames, :]
        np.save(save_path, content_feature.cpu().detach().numpy())


class WhisperExtractor(BaseExtractor):
    def __init__(self, config):
        super(WhisperExtractor, self).__init__(config)
        self.extractor_type = "whisper"

    def load_model(self):
        # load whisper checkpoint
        print("Loading Whisper Model...")

        checkpoint_file = (
            self.cfg.preprocess.whisper_model_path
            if "whisper_model_path" in self.cfg.preprocess
            else None
        )
        model = whisper.load_model(
            self.cfg.preprocess.whisper_model, checkpoint_file=checkpoint_file
        )
        if torch.cuda.is_available():
            print("Using GPU...\n")
            model = model.cuda()
        else:
            print("Using CPU...\n")

        self.model = model.eval()

    def extract_content_features(self, wavs, lens):
        """extract content features from a batch of dataloader
        Args:
            wavs: tensor (batch_size, T)
            lens: list
        """
        # wavs: (batch, max_len)
        wavs = whisper.pad_or_trim(wavs)
        # batch_mel: (batch, 80, 3000)
        batch_mel = whisper.log_mel_spectrogram(wavs).to(self.model.device)
        with torch.no_grad():
            # (batch, 1500, 1024)
            features = self.model.embed_audio(batch_mel)
        return features


class ContentvecExtractor(BaseExtractor):
    def __init__(self, cfg):
        super(ContentvecExtractor, self).__init__(cfg)
        self.extractor_type = "contentvec"

    def load_model(self):
        assert self.model == None
        # Load model
        ckpt_path = self.cfg.preprocess.contentvec_file
        print("Load Contentvec Model...")

        models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
            [ckpt_path],
            suffix="",
        )
        model = models[0]
        model.eval()

        if torch.cuda.is_available():
            # print("Using GPU...\n")
            model = model.cuda()

        self.model = model

    def extract_content_features(self, wavs, lens):
        """extract content features from a batch of dataloader
        Args:
            wavs: tensor (batch, T)
            lens: list
        """
        device = next(self.model.parameters()).device
        wavs = wavs.to(device)  # (batch, max_len)
        padding_mask = torch.eq(wavs, torch.zeros_like(wavs)).to(device)
        with torch.no_grad():
            logits = self.model.extract_features(
                source=wavs, padding_mask=padding_mask, output_layer=12
            )
            # feats: (batch, T, 256)
            feats = self.model.final_proj(logits[0])
        return feats


class WenetExtractor(BaseExtractor):
    def __init__(self, config):
        super(WenetExtractor, self).__init__(config)
        self.extractor_type = "wenet"

    def load_model(self):
        wenet_cfg = self.cfg.preprocess.wenet_config
        wenet_model_path = self.cfg.preprocess.wenet_model_path
        # load Wenet config
        with open(wenet_cfg, "r") as w:
            wenet_configs = yaml.load(w, Loader=yaml.FullLoader)
        self.extract_conf = copy.deepcopy(wenet_configs["dataset_conf"])
        print("Loading Wenet Model...")
        self.model = init_model(wenet_configs)
        load_checkpoint(self.model, wenet_model_path)

        if torch.cuda.is_available():
            print("Using GPU...\n")
            self.model = self.model.cuda()
        else:
            print("Using CPU...\n")

        self.model = self.model.eval()

    def extract_content_features(self, wavs, lens):
        """extract content features from a batch of dataloader
        Args:
            wavs: tensor
            lens: list
        """
        feats_list = []
        lengths_list = []

        device = next(self.model.parameters()).device
        # Extract fbank/mfcc features by kaldi
        assert self.extract_conf is not None, "load model first!"
        feats_type = self.extract_conf.get("feats_type", "fbank")
        assert feats_type in ["fbank", "mfcc"]

        for idx, wav in enumerate(wavs):
            # wav: (T)
            wav = wav[: lens[idx]].to(device)

            # pad one frame to compensate for the frame cut off after feature extraction
            pad_tensor = torch.zeros(160, device=wav.device)
            wav = torch.cat((wav, pad_tensor), dim=-1)
            wav *= 1 << 15

            wav = wav.unsqueeze(0)  # (T) -> (1, T)
            if feats_type == "fbank":
                fbank_conf = self.extract_conf.get("fbank_conf", {})
                feat = kaldi.fbank(
                    wav,
                    sample_frequency=16000,
                    num_mel_bins=fbank_conf["num_mel_bins"],
                    frame_length=fbank_conf["frame_length"],
                    frame_shift=fbank_conf["frame_shift"],
                    dither=fbank_conf["dither"],
                )
            elif feats_type == "mfcc":
                mfcc_conf = self.extract_conf.get("mfcc", {})
                feat = kaldi.mfcc(
                    wav,
                    sample_frequency=16000,
                    num_mel_bins=mfcc_conf["num_mel_bins"],
                    frame_length=mfcc_conf["frame_length"],
                    frame_shift=mfcc_conf["frame_shift"],
                    dither=mfcc_conf["dither"],
                    num_ceps=mfcc_conf.get("num_ceps", 40),
                    high_freq=mfcc_conf.get("high_freq", 0.0),
                    low_freq=mfcc_conf.get("low_freq", 20.0),
                )
            feats_list.append(feat)
            lengths_list.append(feat.shape[0])

        feats_lengths = torch.tensor(lengths_list, dtype=torch.int32).to(device)
        feats_tensor = pad_sequence(feats_list, batch_first=True).to(
            device
        )  # (batch, len, 80)

        features = self.model.encoder_extractor(
            feats_tensor,
            feats_lengths,
            decoding_chunk_size=-1,
            num_decoding_left_chunks=-1,
            simulate_streaming=False,
        )
        return features


class MertExtractor(BaseExtractor):
    def __init__(self, cfg):
        super(MertExtractor, self).__init__(cfg)
        self.extractor_type = "mert"
        self.preprocessor = None

    def load_model(self):
        assert self.model == None
        assert self.preprocessor == None

        print("Loading MERT Model: ...", self.cfg.preprocess.mert_model)

        local_mert_path = "/mnt/workspace/fangzihao/acce/Amphion/pretrained/MERT"

        model_name = self.cfg.preprocess.mert_model
        model = AutoModel.from_pretrained(local_mert_path, trust_remote_code=True)

        if torch.cuda.is_available():
            model = model.cuda()
        preprocessor = Wav2Vec2FeatureExtractor.from_pretrained(
            local_mert_path, trust_remote_code=True
        )

        self.model = model
        self.preprocessor = preprocessor

    def extract_content_features(self, wavs, lens):
        """extract content features from a batch of dataloader
        Args:
            wavs: tensor (batch, T)
            lens: list
        """
        with torch.no_grad():
            sample_rate = self.preprocessor.sampling_rate
            device = next(self.model.parameters()).device
            assert (
                sample_rate == self.cfg.preprocess.mert_sample_rate
            ), "mert sample rate mismatch, expected {}, got {}".format(
                self.cfg.preprocess.mert_sample_rate, sample_rate
            )
            mert_features = []
            # wav: (len)
            for wav in wavs:
                # {input_values: tensor, attention_mask: tensor}
                inputs = self.preprocessor(
                    wavs, sampling_rate=sample_rate, return_tensors="pt"
                ).to(device)

                outputs = self.model(**inputs, output_hidden_states=True)
                # (25 layers, time steps, 1024 feature_dim)
                all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()
                # (1, frame_len, 1024) -> (frame_len, 1024)
                feature = outputs.hidden_states[
                    self.cfg.preprocess.mert_feature_layer
                ].squeeze(0)
                mert_features.append(feature)

        return mert_features


def extract_utt_content_features_dataloader(cfg, metadata, num_workers):
    dataset_name = metadata[0]["Dataset"]

    if cfg.preprocess.extract_whisper_feature:
        feat_dir = os.path.join(cfg.preprocess.processed_dir, dataset_name, "whisper")
        os.makedirs(feat_dir, exist_ok=True)
        feat_files_num = len(os.listdir(feat_dir))

        if feat_files_num != len(metadata):
            whisper_waveforms = FFmpegDataset(
                cfg, dataset_name, cfg.preprocess.whisper_sample_rate, metadata=metadata
            )
            data_loader = DataLoader(
                whisper_waveforms,
                num_workers=num_workers,
                shuffle=False,
                pin_memory=cfg.preprocess.pin_memory,
                batch_size=cfg.preprocess.content_feature_batch_size,
                collate_fn=collate_batch,
                drop_last=False,
            )
            extractor = WhisperExtractor(cfg)
            extractor.load_model()
            for batch_idx, items in enumerate(tqdm(data_loader)):
                _metadata, wavs, lens = items

                batch_content_features = extractor.extract_content_features(
                    wavs,
                    lens,
                )
                for index, utt in enumerate(_metadata):
                    extractor.save_feature(utt, batch_content_features[index])

    if cfg.preprocess.extract_contentvec_feature:
        feat_dir = os.path.join(
            cfg.preprocess.processed_dir, dataset_name, "contentvec"
        )
        os.makedirs(feat_dir, exist_ok=True)
        feat_files_num = len(os.listdir(feat_dir))

        if feat_files_num != len(metadata):
            contentvec_waveforms = LibrosaDataset(
                cfg,
                dataset_name,
                cfg.preprocess.contentvec_sample_rate,
                metadata=metadata,
            )
            data_loader = DataLoader(
                contentvec_waveforms,
                num_workers=num_workers,
                shuffle=False,
                pin_memory=cfg.preprocess.pin_memory,
                batch_size=cfg.preprocess.content_feature_batch_size,
                collate_fn=collate_batch,
                drop_last=False,
            )
            extractor = ContentvecExtractor(cfg)
            extractor.load_model()
            for batch_idx, items in enumerate(tqdm(data_loader)):
                _metadata, wavs, lens = items

                batch_content_features = extractor.extract_content_features(wavs, lens)
                for index, utt in enumerate(_metadata):
                    extractor.save_feature(utt, batch_content_features[index])

    if cfg.preprocess.extract_wenet_feature:
        feat_dir = os.path.join(cfg.preprocess.processed_dir, dataset_name, "wenet")
        os.makedirs(feat_dir, exist_ok=True)
        feat_files_num = len(os.listdir(feat_dir))

        if feat_files_num != len(metadata):
            wenet_waveforms = TorchaudioDataset(
                cfg, dataset_name, cfg.preprocess.wenet_sample_rate, metadata=metadata
            )
            data_loader = DataLoader(
                wenet_waveforms,
                num_workers=num_workers,
                shuffle=False,
                pin_memory=cfg.preprocess.pin_memory,
                batch_size=cfg.preprocess.content_feature_batch_size,
                collate_fn=collate_batch,
                drop_last=False,
            )
            extractor = WenetExtractor(cfg)
            extractor.load_model()
            for batch_idx, items in enumerate(tqdm(data_loader)):
                _metadata, wavs, lens = items

                batch_content_features = extractor.extract_content_features(
                    wavs,
                    lens,
                )
                for index, utt in enumerate(_metadata):
                    extractor.save_feature(utt, batch_content_features[index])

    if cfg.preprocess.extract_mert_feature:
        feat_dir = os.path.join(cfg.preprocess.processed_dir, dataset_name, "mert")
        os.makedirs(feat_dir, exist_ok=True)
        feat_files_num = len(os.listdir(feat_dir))

        if feat_files_num != len(metadata):
            mert_waveforms = TorchaudioDataset(
                cfg, dataset_name, cfg.preprocess.mert_sample_rate, metadata=metadata
            )
            data_loader = DataLoader(
                mert_waveforms,
                num_workers=num_workers,
                shuffle=False,
                pin_memory=cfg.preprocess.pin_memory,
                batch_size=cfg.preprocess.content_feature_batch_size,
                collate_fn=collate_batch,
                drop_last=False,
            )
            extractor = MertExtractor(cfg)
            extractor.load_model()
            for batch_idx, items in enumerate(tqdm(data_loader)):
                _metadata, wavs, lens = items

                batch_content_features = extractor.extract_content_features(
                    wavs,
                    lens,
                )
                for index, utt in enumerate(_metadata):
                    extractor.save_feature(utt, batch_content_features[index])

    if cfg.preprocess.extract_phoneme:
        phone_symbols = set()
        out_path = os.path.join(cfg.preprocess.processed_dir, dataset_name, cfg.preprocess.phone_dir)
        os.makedirs(out_path, exist_ok=True)
        
        if cfg.preprocess.text_extractor in ["espeak", "pypinyin", "pypinyin_initials_finals"]:
            text_tokenizer = G2PModule(backend=cfg.preprocess.text_extractor)
        elif cfg.preprocess.text_extractor == 'lexicon':
            assert cfg.lexicon_path != ""
            lexicon = read_lexicon(cfg.lexicon_path)
        else:
            print('No suppert to', cfg.preprocess.text_extractor)
            raise

        for utt in tqdm(metadata):
            uid = utt["Uid"]
            text = utt["Text"]            
            if cfg.preprocess.text_extractor in ["espeak", "pypinyin", "pypinyin_initials_finals"]:
                text = text.replace("”", '"').replace("“", '"')
                phone = tokenize_text(text_tokenizer, text=text)        
                phone_symbols.update(phone)    
                phone_seq = [phn for phn in phone]
                
            elif cfg.preprocess.text_extractor == 'lexicon':
                phone = preprocess_english(text, lexicon)
                phone_seq = [phn for phn in phone]
               
            phone_path = os.path.join(out_path, uid+'.phone')
            with open(phone_path, 'w') as fin:
                fin.write(' '.join(phone_seq))

        phone_symbols_file = os.path.join(cfg.preprocess.processed_dir, 
                                          dataset_name, 
                                          cfg.preprocess.symbols_dict)


        # load and merge saved phone symbols                
        if os.path.exists(phone_symbols_file):
            phone_symbol_dict_saved = SymbolTable.from_file(phone_symbols_file)._sym2id.keys()
            phone_symbols.update(set(phone_symbol_dict_saved))

        phone_symbol_dict = SymbolTable()
        for s in sorted(list(phone_symbols)):
            phone_symbol_dict.add(s)
        phone_symbol_dict.to_file(phone_symbols_file)            