1"""
  2# Export module
  3
  4This module lets you export data produced with MSA explorer.
  5
  6## Functions:
  7"""
  8
  9import os
 10from numpy import ndarray
 11from msaexplorer import config
 12
 13
 14def _check_and_create_path(path: str):
 15    """
 16    Check and create path if it doesn't exist.
 17    :param path: string to file
 18    """
 19    if path is not None:
 20        output_dir = os.path.dirname(path)
 21        if output_dir and not os.path.exists(output_dir):
 22            os.makedirs(output_dir)
 23
 24
 25def snps(snp_dict: dict, format_type: str = 'vcf', path: str | None = None) -> str | None | ValueError:
 26    """
 27    Export a SNP dictionary to a VCF or tabular format. Importantly, the input dictionary has to be in the standard
 28    format that MSAexplorer produces.
 29
 30    :param snp_dict: Dictionary containing SNP positions and variant information.
 31    :param format_type: Format type ('vcf' or 'tabular'). Default is 'vcf'.
 32    :param path: Path to output VCF or tabular format. (optional)
 33    :return: A string containing the SNP data in the requested format.
 34    :raises ValueError: if the input dictionary is missing required keys or format_type is invalid.
 35    """
 36
 37    def _validate():
 38        if not isinstance(snp_dict, dict):
 39            raise ValueError('Input SNP data must be a dictionary.')
 40        for key in ['#CHROM', 'POS']:
 41            if key not in snp_dict:
 42                raise ValueError(f"Missing required key '{key}' in SNP data.")
 43        if not isinstance(snp_dict['POS'], dict):
 44            raise ValueError('Expected the \'POS\' key to contain a dictionary of positions.')
 45        if format_type not in ['vcf', 'tabular']:
 46            raise ValueError('Invalid format_type.')
 47        _check_and_create_path(path)
 48
 49    def _vcf_format(snp_dict: dict) -> list:
 50        """
 51        Produce  vcf formatted SNP data.
 52        :param snp_dict: dictionary containing SNP positions and variant information.
 53        :return: list of lines to write
 54        """
 55        output_lines = []
 56        # VCF header
 57        output_lines.append('##fileformat=VCFv4.2')
 58        output_lines.append('##source=MSAexplorer')
 59        output_lines.append('#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO')
 60        # process each SNP position in sorted order
 61        for pos in sorted(snp_dict['POS'].keys()):
 62            pos_info = snp_dict['POS'][pos]
 63            ref = pos_info.get('ref', '.')
 64            alt_dict = pos_info.get('ALT', {})
 65            # Create comma-separated list of alternative alleles
 66            alt_alleles = ",".join(alt_dict.keys()) if alt_dict else "."
 67            # Prepare INFO field: include allele frequencies and sequence IDs
 68            afs = []
 69            seq_ids = []
 70            for alt, details in alt_dict.items():
 71                af = details.get('AF', 0)
 72                afs.append(str(af))
 73                seq_ids.append("|".join(details.get('SEQ_ID', [])))
 74            info_fields = []
 75            if afs:
 76                info_fields.append("AF=" + ",".join(afs))
 77            if seq_ids:
 78                info_fields.append("SEQ_ID=" + ",".join(seq_ids))
 79            info = ";".join(info_fields) if info_fields else "."
 80
 81            # VCF is 1-indexed; we assume pos is 0-indexed and add 1
 82            line = f"{snp_dict['#CHROM']}\t{pos + 1}\t.\t{ref}\t{alt_alleles}\t.\t.\t{info}"
 83            output_lines.append(line)
 84
 85        return output_lines
 86
 87    def _tabular_format(snp_dict: dict) -> list:
 88        """
 89        Produce  tabular formatted SNP data.
 90
 91        :param snp_dict: dictionary containing SNP positions and variant information.
 92        :return: list of lines to write
 93        """
 94        output_lines = []
 95        # Create a header for the tabular output
 96        output_lines.append('CHROM\tPOS\tREF\tALT\tAF\tSEQ_ID')
 97
 98        # Process each SNP position and each alternative allele
 99        for pos in sorted(snp_dict['POS'].keys()):
100            pos_info = snp_dict['POS'][pos]
101            ref = pos_info.get('ref', '.')
102            alt_dict = pos_info.get('ALT', {})
103            for alt, details in alt_dict.items():
104                af = details.get('AF', 0)
105                seq_id = ",".join(details.get('SEQ_ID', []))
106                output_lines.append(f"{snp_dict['#CHROM']}\t{pos + 1}\t{ref}\t{alt}\t{af}\t{seq_id}")
107
108        return output_lines
109
110    # validate correct input format
111    _validate()
112
113    # generate line data
114    if format_type == 'vcf':
115        lines = _vcf_format(snp_dict)
116    else:
117        lines = _tabular_format(snp_dict)
118
119    # export to file or return plain text
120    if path is not None:
121        out_path = f"{path}.{format_type}"
122        with open(out_path, 'w') as out_file:
123            out_file.write('\n'.join(lines))
124    else:
125        return '\n'.join(lines)
126
127
128def fasta(sequence: str | dict, header: str | None = None, path: str | None = None) -> str | None:
129    """
130    Export a fasta sequence from str or alignment in dictionary format to either a string or save directly to file.
131    The alignment format must have headers as keys and the corresponding sequence as values.
132    :param sequence: sequence to export
133    :param header: optional header file
134    :param path: path to save the file
135    :return: fasta formatted string
136    """
137    def _validate_sequence(seq: str):
138        if not set(seq).issubset(set(config.POSSIBLE_CHARS)):
139            raise ValueError(f'Sequence contains invalid characters. Detected chars: {set(seq)}')
140
141    _check_and_create_path(path)
142    fasta_formated_sequence = ''
143
144    if type(sequence) is str:
145        _validate_sequence(sequence)
146        fasta_formated_sequence = f'>{header}\n{sequence}'
147    elif type(sequence) is dict:
148        for header, sequence in sequence.items():
149            if type(sequence) is not str:
150                raise ValueError(f'Sequences in the dictionary must be strings.')
151            _validate_sequence(sequence)
152            fasta_formated_sequence = f'{fasta_formated_sequence}\n>{header}\n{sequence}' if fasta_formated_sequence != '' else f'>{header}\n{sequence}'
153
154    if path is not None:
155        with open(path, 'w') as out_file:
156            out_file.write(fasta_formated_sequence)
157    else:
158        return fasta_formated_sequence
159
160
161def stats(stat_data: list | ndarray, seperator: str = '\t', path: str | None = None) -> str | None:
162    """
163    Export a list of stats per nucleotide to tabular or csv format.
164
165    :param stat_data: list of stat values
166    :param seperator: seperator for values and index
167    :param path: path to save the file
168    :return: tabular/csv formatted string
169    """
170    # ini
171    _check_and_create_path(path)
172
173    lines = [f'position{seperator}value']
174
175    for idx, stat_val in enumerate(stat_data):
176        lines.append(f'{idx}{seperator}{stat_val}')
177
178    if path is not None:
179        with open(path, 'w') as out_file:
180            out_file.write('\n'.join(lines))
181    else:
182        return '\n'.join(lines)
183
184
185def orf(orf_dict: dict, chrom: str, path: str | None = None) -> str | ValueError:
186    """
187    Exports the ORF dictionary to a .bed file.
188
189    :param orf_dict: Dictionary containing ORF information.
190    :param path: Path to the output .bed file.
191    :param : Reference name
192    """
193    if not orf_dict:
194        raise ValueError("The ORF dictionary is empty. Nothing to export.")
195    else:
196        if list(orf_dict[list(orf_dict.keys())[0]].keys()) != ['location', 'frame', 'strand', 'conservation', 'internal']:
197            raise ValueError("The ORF dictionary has not the right format.")
198
199    _check_and_create_path(path)
200
201    lines = []
202
203    for orf_id, orf_data in orf_dict.items():
204        lines.append(
205            f"{chrom}\t{orf_data['location'][0][0]}\t{orf_data['location'][0][1]}\t{orf_id}\t{orf_data['conservation']:.2f}\t{orf_data['strand']}"
206        )
207
208    if path is not None:
209        with open(path, 'w') as out_file:
210            out_file.write('\n'.join(lines))
211    else:
212        return '\n'.join(lines)
213
214
215def character_freq(char_dict: dict, seperator: str = '\t', path: str | None = None) -> str | None | ValueError:
216    """
217    Export a character frequency dictionary to tabular or csv format.
218
219    :param char_dict: Dictionary containing the character frequencies.
220    :param seperator: seperator for the table e.g. tabular or comma
221    :param path: Path to output table.
222
223    :return: A string containing the character frequency table.
224    :raises ValueError: if the input dictionary is missing required keys or format_type is invalid.
225    """
226
227    def _validate():
228        if not isinstance(char_dict, dict):
229            raise ValueError('Data must be a dictionary.')
230        for key, value in char_dict.items():
231            for key_2, value_2 in value.items():
232                if key_2 not in config.POSSIBLE_CHARS:
233                    raise ValueError(f'The key {key_2} is invalid.')
234                for key_3, value_3 in value_2.items():
235                    if key_3 not in ['counts', '% of alignment', '% of non-gapped']:
236                        raise ValueError(f'The key "{key_3}" is invalid.')
237
238    # validate input
239    _validate()
240
241    lines = [F'sequence{seperator}char{seperator}counts{seperator}% of non-gapped']
242    for key, value in char_dict.items():
243        if key == 'total':
244            continue
245        for key_2, value_2 in value.items():
246            if key_2 == '-':
247                continue
248            lines.append(f'{key}{seperator}{key_2}{seperator}{value_2["counts"]}{seperator}{value_2["% of non-gapped"]}')
249
250    # export data
251    if path is not None:
252        with open(path, 'w') as out_file:
253            out_file.write('\n'.join(lines))
254    else:
255        return '\n'.join(lines)
256
257
258def percent_recovery(rec_dict: dict, seperator: str = '\t', path: str | None = None) -> str | None | ValueError:
259    """
260    Export percent_recovery dictionary to tabular or csv format.
261
262    :param rec_dict: Dictionary containing the character frequencies.
263    :param seperator: seperator for the table e.g. tabular or comma
264    :param path: Path to output table.
265
266    :return: A string containing the character frequency table.
267    :raises ValueError: if the input dictionary is missing required keys or format_type is invalid.
268    """
269    def _validate():
270        if not isinstance(rec_dict, dict):
271            raise ValueError('Data must be a dictionary.')
272        for key, value in rec_dict.items():
273            if type(key) != str:
274                raise ValueError(f'The key {key} is invalid.')
275            elif type(value) != float:
276                raise ValueError(f'The value {value} is invalid.')
277
278    # validate input
279    _validate()
280
281    lines = [F'sequence{seperator}% recovery']
282    for key, value in rec_dict.items():
283        lines.append(
284            f'{key}{seperator}{value}')
285
286    # export data
287    if path is not None:
288        with open(path, 'w') as out_file:
289            out_file.write('\n'.join(lines))
290    else:
291        return '\n'.join(lines)