1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Code to write out hyperdrive source lists.
//!
//! To make the source list files a little easier to read and write,
//! [SourceList] isn't directly serialisable or deserialisable. Use temporary
//! types here to do the serde magic and write out a [SourceList].

use std::collections::BTreeMap;

use log::warn;

use super::{TmpComponent, TmpFluxDensityType, TmpSourceList};
use crate::srclist::{error::WriteSourceListError, ComponentType, FluxDensityType, SourceList};

fn source_list_to_tmp_sl(sl: &SourceList, num_sources: Option<usize>) -> TmpSourceList {
    let mut tmp_sl: BTreeMap<String, Vec<TmpComponent>> = BTreeMap::new();

    let mut num_written_sources = 0;
    // Note that, if sorted, each source in the source list is dimmer than the
    // last!
    for (name, source) in sl.iter() {
        // If `num_sources` is supplied, then check that we're not writing out
        // too many sources.
        if let Some(num_sources) = num_sources {
            if num_written_sources == num_sources {
                break;
            }
        }

        let mut tmp_comps = Vec::with_capacity(source.components.len());
        for comp in &source.components {
            let comp_type = match &comp.comp_type {
                ComponentType::Point => ComponentType::Point,
                ComponentType::Gaussian { maj, min, pa } => ComponentType::Gaussian {
                    maj: maj.to_degrees() * 3600.0,
                    min: min.to_degrees() * 3600.0,
                    pa: pa.to_degrees(),
                },
                ComponentType::Shapelet {
                    maj,
                    min,
                    pa,
                    coeffs,
                } => ComponentType::Shapelet {
                    maj: maj.to_degrees() * 3600.0,
                    min: min.to_degrees() * 3600.0,
                    pa: pa.to_degrees(),
                    coeffs: coeffs.clone(),
                },
            };

            tmp_comps.push(TmpComponent {
                ra: comp.radec.ra.to_degrees(),
                dec: comp.radec.dec.to_degrees(),
                comp_type,
                flux_type: match &comp.flux_type {
                    FluxDensityType::List { fds } => TmpFluxDensityType::List(fds.to_vec()),
                    FluxDensityType::PowerLaw { si, fd } => TmpFluxDensityType::PowerLaw {
                        si: *si,
                        fd: fd.clone(),
                    },
                    FluxDensityType::CurvedPowerLaw { si, fd, q } => {
                        TmpFluxDensityType::CurvedPowerLaw {
                            si: *si,
                            fd: fd.clone(),
                            q: *q,
                        }
                    }
                },
            })
        }

        tmp_sl.insert(name.clone(), tmp_comps);
        num_written_sources += 1;
    }

    if let Some(num_sources) = num_sources {
        if num_sources > num_written_sources {
            warn!("Couldn't write the requested number of sources ({num_sources}): wrote {num_written_sources}")
        }
    }

    tmp_sl
}

/// Write a `SourceList` to a yaml file.
pub(crate) fn source_list_to_yaml<T: std::io::Write>(
    buf: &mut T,
    sl: &SourceList,
    num_sources: Option<usize>,
) -> Result<(), WriteSourceListError> {
    let tmp_sl = source_list_to_tmp_sl(sl, num_sources);
    serde_yaml::to_writer(buf, &tmp_sl)?;
    Ok(())
}

/// Write a `SourceList` to a json file.
pub(crate) fn source_list_to_json<T: std::io::Write>(
    buf: &mut T,
    sl: &SourceList,
    num_sources: Option<usize>,
) -> Result<(), WriteSourceListError> {
    let tmp_sl = source_list_to_tmp_sl(sl, num_sources);
    serde_json::to_writer_pretty(buf, &tmp_sl)?;
    Ok(())
}