text: # -*- coding: utf-8 -*-
  # Copyright 2020 TensorFlowTTS Team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Perform preprocessing and raw feature extraction for Baker dataset."""
  
  import os
  import re
  from typing import Dict, List, Union, Tuple, Any
  
  import librosa
  import numpy as np
  import soundfile as sf
  from dataclasses import dataclass, field
  from pypinyin import Style
  from pypinyin.contrib.neutral_tone import NeutralToneWith5Mixin
  from pypinyin.converter import DefaultConverter
  from pypinyin.core import Pinyin
  from tensorflow_tts.processor import BaseProcessor
  from tensorflow_tts.utils.utils import PROCESSOR_FILE_NAME
  
  _pad = ["pad"]
  _eos = ["eos"]
  _pause = ["sil", "#0", "#1", "#2", "#3"]
  
  _initials = ["^", "b", "c", "ch", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "sh", "t", "x", "z", "zh", ]
  
  _tones = ["1", "2", "3", "4", "5"]
  
  _finals = ["a", "ai", "an", "ang", "ao", "e", "ei", "en", "eng", "er", "i", "ia", "ian", "iang", "iao", "ie", "ii", "iii", "in", "ing", "iong", "iou", "o", "ong", "ou", "u", "ua", "uai", "uan",
      "uang", "uei", "uen", "ueng", "uo", "v", "van", "ve", "vn", ]
  
  BAKER_SYMBOLS = _pad + _pause + _initials + [i + j for i in _finals for j in _tones] + _eos
  
  PINYIN_DICT = {"a": ("^", "a"), "ai": ("^", "ai"), "an": ("^", "an"), "ang": ("^", "ang"), "ao": ("^", "ao"), "ba": ("b", "a"), "bai": ("b", "ai"), "ban": ("b", "an"), "bang": ("b", "ang"),
      "bao": ("b", "ao"), "be": ("b", "e"), "bei": ("b", "ei"), "ben": ("b", "en"), "beng": ("b", "eng"), "bi": ("b", "i"), "bian": ("b", "ian"), "biao": ("b", "iao"), "bie": ("b", "ie"),
      "bin": ("b", "in"), "bing": ("b", "ing"), "bo": ("b", "o"), "bu": ("b", "u"), "ca": ("c", "a"), "cai": ("c", "ai"), "can": ("c", "an"), "cang": ("c", "ang"), "cao": ("c", "ao"), "ce": ("c", "e"),
      "cen": ("c", "en"), "ceng": ("c", "eng"), "cha": ("ch", "a"), "chai": ("ch", "ai"), "chan": ("ch", "an"), "chang": ("ch", "ang"), "chao": ("ch", "ao"), "che": ("ch", "e"), "chen": ("ch", "en"),
      "cheng": ("ch", "eng"), "chi": ("ch", "iii"), "chong": ("ch", "ong"), "chou": ("ch", "ou"), "chu": ("ch", "u"), "chua": ("ch", "ua"), "chuai": ("ch", "uai"), "chuan": ("ch", "uan"),
      "chuang": ("ch", "uang"), "chui": ("ch", "uei"), "chun": ("ch", "uen"), "chuo": ("ch", "uo"), "ci": ("c", "ii"), "cong": ("c", "ong"), "cou": ("c", "ou"), "cu": ("c", "u"), "cuan": ("c", "uan"),
      "cui": ("c", "uei"), "cun": ("c", "uen"), "cuo": ("c", "uo"), "da": ("d", "a"), "dai": ("d", "ai"), "dan": ("d", "an"), "dang": ("d", "ang"), "dao": ("d", "ao"), "de": ("d", "e"),
      "dei": ("d", "ei"), "den": ("d", "en"), "deng": ("d", "eng"), "di": ("d", "i"), "dia": ("d", "ia"), "dian": ("d", "ian"), "diao": ("d", "iao"), "die": ("d", "ie"), "ding": ("d", "ing"),
      "diu": ("d", "iou"), "dong": ("d", "ong"), "dou": ("d", "ou"), "du": ("d", "u"), "duan": ("d", "uan"), "dui": ("d", "uei"), "dun": ("d", "uen"), "duo": ("d", "uo"), "e": ("^", "e"),
      "ei": ("^", "ei"), "en": ("^", "en"), "ng": ("^", "en"), "eng": ("^", "eng"), "er": ("^", "er"), "fa": ("f", "a"), "fan": ("f", "an"), "fang": ("f", "ang"), "fei": ("f", "ei"), "fen": ("f", "en"),
      "feng": ("f", "eng"), "fo": ("f", "o"), "fou": ("f", "ou"), "fu": ("f", "u"), "ga": ("g", "a"), "gai": ("g", "ai"), "gan": ("g", "an"), "gang": ("g", "ang"), "gao": ("g", "ao"), "ge": ("g", "e"),
      "gei": ("g", "ei"), "gen": ("g", "en"), "geng": ("g", "eng"), "gong": ("g", "ong"), "gou": ("g", "ou"), "gu": ("g", "u"), "gua": ("g", "ua"), "guai": ("g", "uai"), "guan": ("g", "uan"),
      "guang": ("g", "uang"), "gui": ("g", "uei"), "gun": ("g", "uen"), "guo": ("g", "uo"), "ha": ("h", "a"), "hai": ("h", "ai"), "han": ("h", "an"), "hang": ("h", "ang"), "hao": ("h", "ao"),
      "he": ("h", "e"), "hei": ("h", "ei"), "hen": ("h", "en"), "heng": ("h", "eng"), "hong": ("h", "ong"), "hou": ("h", "ou"), "hu": ("h", "u"), "hua": ("h", "ua"), "huai": ("h", "uai"),
      "huan": ("h", "uan"), "huang": ("h", "uang"), "hui": ("h", "uei"), "hun": ("h", "uen"), "huo": ("h", "uo"), "ji": ("j", "i"), "jia": ("j", "ia"), "jian": ("j", "ian"), "jiang": ("j", "iang"),
      "jiao": ("j", "iao"), "jie": ("j", "ie"), "jin": ("j", "in"), "jing": ("j", "ing"), "jiong": ("j", "iong"), "jiu": ("j", "iou"), "ju": ("j", "v"), "juan": ("j", "van"), "jue": ("j", "ve"),
      "jun": ("j", "vn"), "ka": ("k", "a"), "kai": ("k", "ai"), "kan": ("k", "an"), "kang": ("k", "ang"), "kao": ("k", "ao"), "ke": ("k", "e"), "kei": ("k", "ei"), "ken": ("k", "en"),
      "keng": ("k", "eng"), "kong": ("k", "ong"), "kou": ("k", "ou"), "ku": ("k", "u"), "kua": ("k", "ua"), "kuai": ("k", "uai"), "kuan": ("k", "uan"), "kuang": ("k", "uang"), "kui": ("k", "uei"),
      "kun": ("k", "uen"), "kuo": ("k", "uo"), "la": ("l", "a"), "lai": ("l", "ai"), "lan": ("l", "an"), "lang": ("l", "ang"), "lao": ("l", "ao"), "le": ("l", "e"), "lei": ("l", "ei"),
      "leng": ("l", "eng"), "li": ("l", "i"), "lia": ("l", "ia"), "lian": ("l", "ian"), "liang": ("l", "iang"), "liao": ("l", "iao"), "lie": ("l", "ie"), "lin": ("l", "in"), "ling": ("l", "ing"),
      "liu": ("l", "iou"), "lo": ("l", "o"), "long": ("l", "ong"), "lou": ("l", "ou"), "lu": ("l", "u"), "lv": ("l", "v"), "luan": ("l", "uan"), "lve": ("l", "ve"), "lue": ("l", "ve"),
      "lun": ("l", "uen"), "luo": ("l", "uo"), "ma": ("m", "a"), "mai": ("m", "ai"), "man": ("m", "an"), "mang": ("m", "ang"), "mao": ("m", "ao"), "me": ("m", "e"), "mei": ("m", "ei"),
      "men": ("m", "en"), "meng": ("m", "eng"), "mi": ("m", "i"), "mian": ("m", "ian"), "miao": ("m", "iao"), "mie": ("m", "ie"), "min": ("m", "in"), "ming": ("m", "ing"), "miu": ("m", "iou"),
      "mo": ("m", "o"), "mou": ("m", "ou"), "mu": ("m", "u"), "na": ("n", "a"), "nai": ("n", "ai"), "nan": ("n", "an"), "nang": ("n", "ang"), "nao": ("n", "ao"), "ne": ("n", "e"), "nei": ("n", "ei"),
      "nen": ("n", "en"), "neng": ("n", "eng"), "ni": ("n", "i"), "nia": ("n", "ia"), "nian": ("n", "ian"), "niang": ("n", "iang"), "niao": ("n", "iao"), "nie": ("n", "ie"), "nin": ("n", "in"),
      "ning": ("n", "ing"), "niu": ("n", "iou"), "nong": ("n", "ong"), "nou": ("n", "ou"), "nu": ("n", "u"), "nv": ("n", "v"), "nuan": ("n", "uan"), "nve": ("n", "ve"), "nue": ("n", "ve"),
      "nuo": ("n", "uo"), "o": ("^", "o"), "ou": ("^", "ou"), "pa": ("p", "a"), "pai": ("p", "ai"), "pan": ("p", "an"), "pang": ("p", "ang"), "pao": ("p", "ao"), "pe": ("p", "e"), "pei": ("p", "ei"),
      "pen": ("p", "en"), "peng": ("p", "eng"), "pi": ("p", "i"), "pian": ("p", "ian"), "piao": ("p", "iao"), "pie": ("p", "ie"), "pin": ("p", "in"), "ping": ("p", "ing"), "po": ("p", "o"),
      "pou": ("p", "ou"), "pu": ("p", "u"), "qi": ("q", "i"), "qia": ("q", "ia"), "qian": ("q", "ian"), "qiang": ("q", "iang"), "qiao": ("q", "iao"), "qie": ("q", "ie"), "qin": ("q", "in"),
      "qing": ("q", "ing"), "qiong": ("q", "iong"), "qiu": ("q", "iou"), "qu": ("q", "v"), "quan": ("q", "van"), "que": ("q", "ve"), "qun": ("q", "vn"), "ran": ("r", "an"), "rang": ("r", "ang"),
      "rao": ("r", "ao"), "re": ("r", "e"), "ren": ("r", "en"), "reng": ("r", "eng"), "ri": ("r", "iii"), "rong": ("r", "ong"), "rou": ("r", "ou"), "ru": ("r", "u"), "rua": ("r", "ua"),
      "ruan": ("r", "uan"), "rui": ("r", "uei"), "run": ("r", "uen"), "ruo": ("r", "uo"), "sa": ("s", "a"), "sai": ("s", "ai"), "san": ("s", "an"), "sang": ("s", "ang"), "sao": ("s", "ao"),
      "se": ("s", "e"), "sen": ("s", "en"), "seng": ("s", "eng"), "sha": ("sh", "a"), "shai": ("sh", "ai"), "shan": ("sh", "an"), "shang": ("sh", "ang"), "shao": ("sh", "ao"), "she": ("sh", "e"),
      "shei": ("sh", "ei"), "shen": ("sh", "en"), "sheng": ("sh", "eng"), "shi": ("sh", "iii"), "shou": ("sh", "ou"), "shu": ("sh", "u"), "shua": ("sh", "ua"), "shuai": ("sh", "uai"),
      "shuan": ("sh", "uan"), "shuang": ("sh", "uang"), "shui": ("sh", "uei"), "shun": ("sh", "uen"), "shuo": ("sh", "uo"), "si": ("s", "ii"), "song": ("s", "ong"), "sou": ("s", "ou"), "su": ("s", "u"),
      "suan": ("s", "uan"), "sui": ("s", "uei"), "sun": ("s", "uen"), "suo": ("s", "uo"), "ta": ("t", "a"), "tai": ("t", "ai"), "tan": ("t", "an"), "tang": ("t", "ang"), "tao": ("t", "ao"),
      "te": ("t", "e"), "tei": ("t", "ei"), "teng": ("t", "eng"), "ti": ("t", "i"), "tian": ("t", "ian"), "tiao": ("t", "iao"), "tie": ("t", "ie"), "ting": ("t", "ing"), "tong": ("t", "ong"),
      "tou": ("t", "ou"), "tu": ("t", "u"), "tuan": ("t", "uan"), "tui": ("t", "uei"), "tun": ("t", "uen"), "tuo": ("t", "uo"), "wa": ("^", "ua"), "wai": ("^", "uai"), "wan": ("^", "uan"),
      "wang": ("^", "uang"), "wei": ("^", "uei"), "wen": ("^", "uen"), "weng": ("^", "ueng"), "wo": ("^", "uo"), "wu": ("^", "u"), "xi": ("x", "i"), "xia": ("x", "ia"), "xian": ("x", "ian"),
      "xiang": ("x", "iang"), "xiao": ("x", "iao"), "xie": ("x", "ie"), "xin": ("x", "in"), "xing": ("x", "ing"), "xiong": ("x", "iong"), "xiu": ("x", "iou"), "xu": ("x", "v"), "xuan": ("x", "van"),
      "xue": ("x", "ve"), "xun": ("x", "vn"), "ya": ("^", "ia"), "yan": ("^", "ian"), "yang": ("^", "iang"), "yao": ("^", "iao"), "ye": ("^", "ie"), "yi": ("^", "i"), "yin": ("^", "in"),
      "ying": ("^", "ing"), "yo": ("^", "iou"), "yong": ("^", "iong"), "you": ("^", "iou"), "yu": ("^", "v"), "yuan": ("^", "van"), "yue": ("^", "ve"), "yun": ("^", "vn"), "za": ("z", "a"),
      "zai": ("z", "ai"), "zan": ("z", "an"), "zang": ("z", "ang"), "zao": ("z", "ao"), "ze": ("z", "e"), "zei": ("z", "ei"), "zen": ("z", "en"), "zeng": ("z", "eng"), "zha": ("zh", "a"),
      "zhai": ("zh", "ai"), "zhan": ("zh", "an"), "zhang": ("zh", "ang"), "zhao": ("zh", "ao"), "zhe": ("zh", "e"), "zhei": ("zh", "ei"), "zhen": ("zh", "en"), "zheng": ("zh", "eng"),
      "zhi": ("zh", "iii"), "zhong": ("zh", "ong"), "zhou": ("zh", "ou"), "zhu": ("zh", "u"), "zhua": ("zh", "ua"), "zhuai": ("zh", "uai"), "zhuan": ("zh", "uan"), "zhuang": ("zh", "uang"),
      "zhui": ("zh", "uei"), "zhun": ("zh", "uen"), "zhuo": ("zh", "uo"), "zi": ("z", "ii"), "zong": ("z", "ong"), "zou": ("z", "ou"), "zu": ("z", "u"), "zuan": ("z", "uan"), "zui": ("z", "uei"),
      "zun": ("z", "uen"), "zuo": ("z", "uo"), }
  
  zh_pattern = re.compile("[\u4e00-\u9fa5]")
  
  
  def is_zh(word):
      global zh_pattern
      match = zh_pattern.search(word)
      return match is not None
  
  
  class MyConverter(NeutralToneWith5Mixin, DefaultConverter):
      pass
  
  
  @dataclass
  class BakerProcessor(BaseProcessor):
      pinyin_dict: Dict[str, Tuple[str, str]] = field(default_factory = lambda: PINYIN_DICT)
      cleaner_names: str = None
      target_rate: int = 24000
      speaker_name: str = "baker"
  
      def __post_init__(self):
          super().__post_init__()
          self.pinyin_parser = self.get_pinyin_parser()
  
      def setup_eos_token(self):
          return _eos[0]
  
      def save_pretrained(self, saved_path):
          os.makedirs(saved_path, exist_ok = True)
          self._save_mapper(os.path.join(saved_path, PROCESSOR_FILE_NAME), {"pinyin_dict": self.pinyin_dict}, )
  
      def create_items(self):
          items = []
          if self.data_dir:
              with open(os.path.join(self.data_dir, "ProsodyLabeling/000001-010000.txt"), encoding = "utf-8", ) as ttf:
                  lines = ttf.readlines()
                  for idx in range(0, len(lines), 2):
                      utt_id, chn_char = lines[idx].strip().split()  # 文件名称, 汉字
                      pinyin = lines[idx + 1].strip().split()  # 拼音
                      if "IY1" in pinyin or "Ｂ" in chn_char:  # 标贝的数据中IY1是对应字幕P，B对应字幕B。这里跳过这两句，不进行处理
                          print(f"Skip this: {utt_id} {chn_char} {pinyin}")
                          continue
                      phonemes = self.get_phoneme_from_char_and_pinyin(chn_char, pinyin)
                      wav_path = os.path.join(self.data_dir, "Wave", "%s.wav" % utt_id)
                      # item格式：音素，音频路径，文件名称，baker。["sil k a2 ^ er2 p u3 #2 p ei2 w ai4 s un1 #1 w an2 h ua2 t i1 sil", "/content/..../1.wav", "1", "baker"]
                      items.append([" ".join(phonemes), wav_path, utt_id, self.speaker_name])
              #
              self.items = items
  
      # 输入：(卡尔普#2陪外孙#1玩滑梯#4, ka2 er2 pu3 pei2 wai4 sun1 wan2 hua2 ti1)
      # 输出：[sil, k, a2, ^, er2, p, u3,#2, p, ei2, w,ai4, s,un1, #1, w, an2, h, ua2, t, i1, sil]
      def get_phoneme_from_char_and_pinyin(self, chn_char, pinyin):  # chn_char，中文字符list；pinyin，拼音list
          # we do not need #4, use sil to replace it
          chn_char = chn_char.replace("#4", "")
          char_len = len(chn_char)
          i, j = 0, 0
          result = ["sil"]
          while i < char_len:
              cur_char = chn_char[i]
              if is_zh(cur_char):
                  if pinyin[j][:-1] not in self.pinyin_dict:  # 找出拼音（不包含声调）不在拼音表中的
                      assert chn_char[i + 1] == "儿"
                      assert pinyin[j][-2] == "r"
                      tone = pinyin[j][-1]
                      a = pinyin[j][:-2]
                      a1, a2 = self.pinyin_dict[a]
                      result += [a1, a2 + tone, "er5"]
                      if i + 2 < char_len and chn_char[i + 2] != "#":
                          result.append("#0")
  
                      i += 2
                      j += 1
                  else:
                      tone = pinyin[j][-1]  # 声调
                      a = pinyin[j][:-1]  # 拼音
                      a1, a2 = self.pinyin_dict[a]
                      result += [a1, a2 + tone]  # 例如：[b, iao1]
  
                      if i + 1 < char_len and chn_char[i + 1] != "#":  # 不是最后一个，但是下一个汉字是#开头的，在拼音list中拼上#0
                          result.append("#0")
  
                      i += 1
                      j += 1
              elif cur_char == "#":  # 遇到#，在拼音list上#和后一位
                  result.append(chn_char[i: i + 2])
                  i += 2
              else:
                  # ignore the unknown char and punctuation
                  # result.append(chn_char[i])
                  i += 1
          if result[-1] == "#0":
              result = result[:-1]
          result.append("sil")
          assert j == len(pinyin)
          return result
  
      def get_one_sample(self, item):
          text, wav_file, utt_id, speaker_name = item
  
          # normalize audio signal to be [-1, 1], soundfile already norm.
          audio, rate = sf.read(wav_file)
          audio = audio.astype(np.float32)
          if rate != self.target_rate:
              assert rate > self.target_rate
              audio = librosa.resample(audio, rate, self.target_rate)
  
          # convert text to ids
          try:
              text_ids = np.asarray(self.text_to_sequence(text), np.int32)
          except Exception as e:
              print(e, utt_id, text)
              return None
  
          # return None
          sample = {"raw_text": text,  # 拼音句子
              "text_ids": text_ids,  # 字母的id列表
              "audio": audio,  # 音频数据
              "utt_id": str(int(utt_id)),  # 文件名
              "speaker_name": speaker_name,  # 说话者
              "rate": self.target_rate,  # 采样率
          }
  
          return sample
  
      def get_pinyin_parser(self):
          my_pinyin = Pinyin(MyConverter())
          pinyin = my_pinyin.pinyin
          return pinyin
  
      def text_to_sequence(self, text, inference = False):
          if inference:
              pinyin = self.pinyin_parser(text, style = Style.TONE3, errors = "ignore")
              new_pinyin = []
              for x in pinyin:
                  x = "".join(x)
                  if "#" not in x:
                      new_pinyin.append(x)
              phonemes = self.get_phoneme_from_char_and_pinyin(text, new_pinyin)
              text = " ".join(phonemes)
              print(f"phoneme seq: {text}")
  
          sequence = []
          # text: sil k a2 ^ er2 p u3 #2 p ei2 w ai4 s un1 #1 w an2 h ua2 t i1 sil
          for symbol in text.split():
              idx = self.symbol_to_id[symbol]
              sequence.append(idx)
  
          # add eos tokens
          sequence += [self.eos_id]
          # 字母会生成一个索引表，在mapper.json文件中
          return sequence
  